From 65182fc38ebd0485d25d1412d70bdc2e90669421 Mon Sep 17 00:00:00 2001 From: sakshichitnis27 <156598682+sakshichitnis27@users.noreply.github.com> Date: Sun, 6 Sep 2026 09:01:30 +0000 Subject: [PATCH] [rebalance] Support preferred leader election goal --- .../org/apache/fluss/client/admin/Admin.java | 3 + .../fluss/client/admin/RebalanceITCase.java | 126 ++++++++++++++++ .../fluss/cluster/rebalance/GoalType.java | 10 +- .../fluss/cluster/rebalance/GoalTypeTest.java | 33 +++++ .../flink/procedure/RebalanceProcedure.java | 2 + .../flink/procedure/FlinkProcedureITCase.java | 35 +++++ .../crates/fluss/src/metadata/goal_type.rs | 4 + .../CoordinatorEventProcessor.java | 19 ++- .../rebalance/RebalanceManager.java | 33 ++++- .../rebalance/goal/GoalOptimizer.java | 14 ++ .../coordinator/rebalance/goal/GoalUtils.java | 2 + .../goal/PreferredLeaderElectionGoal.java | 95 ++++++++++++ .../rebalance/model/ClusterModel.java | 11 +- .../rebalance/model/ReplicaModel.java | 14 ++ .../statemachine/ReplicaLeaderElection.java | 9 ++ .../goal/PreferredLeaderElectionGoalTest.java | 138 ++++++++++++++++++ .../ReplicaLeaderElectionTest.java | 4 + .../docs/maintenance/operations/rebalance.md | 17 +++ 18 files changed, 563 insertions(+), 6 deletions(-) create mode 100644 fluss-common/src/test/java/org/apache/fluss/cluster/rebalance/GoalTypeTest.java create mode 100644 fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/PreferredLeaderElectionGoal.java create mode 100644 fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/goal/PreferredLeaderElectionGoalTest.java diff --git a/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java b/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java index 5d749b6d432..07357f0e5d3 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java @@ -661,6 +661,9 @@ ListOffsetsResult listOffsets( *

More details, Fluss collects the cluster's load information and optimizes to perform load * balancing according to the user-defined {@code priorityGoals}. * + *

{@link GoalType#PREFERRED_LEADER_ELECTION} must be requested as a standalone goal. It + * changes only bucket leadership and leaves replica assignments unchanged. + * *

Currently, Fluss only supports one active rebalance task in the cluster. If an uncompleted * rebalance task exists, Fluss will return the uncompleted rebalance task's progress. * diff --git a/fluss-client/src/test/java/org/apache/fluss/client/admin/RebalanceITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/admin/RebalanceITCase.java index 0800afdfddb..26c42f1604b 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/admin/RebalanceITCase.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/admin/RebalanceITCase.java @@ -29,10 +29,12 @@ import org.apache.fluss.exception.RebalanceFailureException; import org.apache.fluss.metadata.DatabaseDescriptor; import org.apache.fluss.metadata.PartitionSpec; +import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableDescriptor; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.server.replica.ReplicaManager; import org.apache.fluss.server.testutils.FlussClusterExtension; +import org.apache.fluss.server.zk.ZooKeeperClient; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -40,8 +42,10 @@ import org.junit.jupiter.api.extension.RegisterExtension; import java.time.Duration; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.List; import java.util.Optional; import static org.apache.fluss.record.TestData.DATA1_SCHEMA; @@ -209,6 +213,128 @@ void testRebalanceForLogTable() throws Exception { admin.removeServerTag(Collections.singletonList(0), ServerTag.PERMANENT_OFFLINE).get(); } + @Test + void testPreferredLeaderElectionAfterRecovery() throws Exception { + String dbName = "db-preferred-leader"; + admin.createDatabase(dbName, DatabaseDescriptor.EMPTY, false).get(); + long tableId = + createTable( + new TablePath(dbName, "preferred-leader-table"), DATA1_TABLE_DESCRIPTOR); + FLUSS_CLUSTER_EXTENSION.waitUntilTableReady(tableId); + + TableBucket tableBucket = new TableBucket(tableId, 0); + ZooKeeperClient zkClient = FLUSS_CLUSTER_EXTENSION.getZooKeeperClient(); + List assignment = + new ArrayList<>( + zkClient.getTableAssignment(tableId) + .get() + .getBucketAssignment(tableBucket.getBucket()) + .getReplicas()); + int preferredLeader = assignment.get(0); + assertThat(FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket)) + .isEqualTo(preferredLeader); + + boolean preferredLeaderStopped = false; + boolean preferredLeaderTagged = false; + try { + FLUSS_CLUSTER_EXTENSION.stopTabletServer(preferredLeader); + preferredLeaderStopped = true; + FLUSS_CLUSTER_EXTENSION.waitUntilReplicaShrinkFromIsr(tableBucket, preferredLeader); + retry( + Duration.ofMinutes(1), + () -> + assertThat(FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket)) + .isNotEqualTo(preferredLeader)); + int failoverLeader = FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket); + + String unavailableRebalanceId = + admin.rebalance(Collections.singletonList(GoalType.PREFERRED_LEADER_ELECTION)) + .get(); + waitUntilRebalanceCompletes(unavailableRebalanceId); + assertThat(FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket)) + .isEqualTo(failoverLeader); + assertThat( + zkClient.getTableAssignment(tableId) + .get() + .getBucketAssignment(tableBucket.getBucket()) + .getReplicas()) + .containsExactlyElementsOf(assignment); + + FLUSS_CLUSTER_EXTENSION.startTabletServer(preferredLeader); + preferredLeaderStopped = false; + FLUSS_CLUSTER_EXTENSION.waitUntilReplicaExpandToIsr(tableBucket, preferredLeader); + + admin.addServerTag( + Collections.singletonList(preferredLeader), ServerTag.TEMPORARY_OFFLINE) + .get(); + preferredLeaderTagged = true; + String taggedRebalanceId = + admin.rebalance(Collections.singletonList(GoalType.PREFERRED_LEADER_ELECTION)) + .get(); + waitUntilRebalanceCompletes(taggedRebalanceId); + assertThat(FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket)) + .isEqualTo(failoverLeader); + + admin.removeServerTag( + Collections.singletonList(preferredLeader), ServerTag.TEMPORARY_OFFLINE) + .get(); + preferredLeaderTagged = false; + + String rebalanceId = + admin.rebalance(Collections.singletonList(GoalType.PREFERRED_LEADER_ELECTION)) + .get(); + waitUntilRebalanceCompletes(rebalanceId); + assertThat(FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket)) + .isEqualTo(preferredLeader); + assertThat( + zkClient.getTableAssignment(tableId) + .get() + .getBucketAssignment(tableBucket.getBucket()) + .getReplicas()) + .containsExactlyElementsOf(assignment); + + int leaderEpoch = zkClient.getLeaderAndIsr(tableBucket).get().leaderEpoch(); + String idempotentRebalanceId = + admin.rebalance(Collections.singletonList(GoalType.PREFERRED_LEADER_ELECTION)) + .get(); + waitUntilRebalanceCompletes(idempotentRebalanceId); + assertThat(zkClient.getLeaderAndIsr(tableBucket).get().leaderEpoch()) + .isEqualTo(leaderEpoch); + + assertThatThrownBy( + () -> + admin.rebalance( + Arrays.asList( + GoalType.PREFERRED_LEADER_ELECTION, + GoalType.LEADER_DISTRIBUTION)) + .get()) + .rootCause() + .isInstanceOf(RebalanceFailureException.class) + .hasMessageContaining("must be used as a standalone rebalance goal"); + } finally { + if (preferredLeaderTagged) { + admin.removeServerTag( + Collections.singletonList(preferredLeader), + ServerTag.TEMPORARY_OFFLINE) + .get(); + } + if (preferredLeaderStopped) { + FLUSS_CLUSTER_EXTENSION.startTabletServer(preferredLeader); + } + } + } + + private void waitUntilRebalanceCompletes(String rebalanceId) { + retry( + Duration.ofMinutes(2), + () -> { + Optional progress = + admin.listRebalanceProgress(rebalanceId).get(); + assertThat(progress).isPresent(); + assertThat(progress.get().status()).isEqualTo(RebalanceStatus.COMPLETED); + }); + } + @Test void testListRebalanceProgress() throws Exception { String dbName = "db-rebalance-list"; diff --git a/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/GoalType.java b/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/GoalType.java index 7a01116c79e..6ac23758500 100644 --- a/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/GoalType.java +++ b/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/GoalType.java @@ -44,7 +44,13 @@ public enum GoalType { * Goal to generate replica movement tasks to ensure that the number of replicas on each * tabletServer is near balanced and the replicas are distributed across racks. */ - RACK_AWARE(2); + RACK_AWARE(2), + + /** + * Goal to move leadership to the first replica in each persisted bucket assignment without + * changing replica assignments. + */ + PREFERRED_LEADER_ELECTION(3); public final int value; @@ -59,6 +65,8 @@ public static GoalType valueOf(int value) { return LEADER_DISTRIBUTION; } else if (value == RACK_AWARE.value) { return RACK_AWARE; + } else if (value == PREFERRED_LEADER_ELECTION.value) { + return PREFERRED_LEADER_ELECTION; } else { throw new IllegalArgumentException( String.format( diff --git a/fluss-common/src/test/java/org/apache/fluss/cluster/rebalance/GoalTypeTest.java b/fluss-common/src/test/java/org/apache/fluss/cluster/rebalance/GoalTypeTest.java new file mode 100644 index 00000000000..9cc551ac22a --- /dev/null +++ b/fluss-common/src/test/java/org/apache/fluss/cluster/rebalance/GoalTypeTest.java @@ -0,0 +1,33 @@ +/* + * 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.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Test for {@link GoalType}. */ +class GoalTypeTest { + + @Test + void testPreferredLeaderElectionFromName() { + assertThat(GoalType.fromName("preferred_leader_election")) + .isEqualTo(GoalType.PREFERRED_LEADER_ELECTION); + assertThat(GoalType.valueOf(3)).isEqualTo(GoalType.PREFERRED_LEADER_ELECTION); + } +} diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/RebalanceProcedure.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/RebalanceProcedure.java index 96772ffe378..8266a4d5466 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/RebalanceProcedure.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/RebalanceProcedure.java @@ -41,6 +41,8 @@ * CALL sys.rebalance('REPLICA_DISTRIBUTION'); * -- Trigger rebalance with REPLICA_DISTRIBUTION and LEADER_DISTRIBUTION goals * CALL sys.rebalance('REPLICA_DISTRIBUTION,LEADER_DISTRIBUTION'); + * -- Restore leadership to the first replica in each assignment + * CALL sys.rebalance('PREFERRED_LEADER_ELECTION'); * */ public class RebalanceProcedure extends ProcedureBase { diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/procedure/FlinkProcedureITCase.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/procedure/FlinkProcedureITCase.java index f1c71f0c26c..4d2e01eeec7 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/procedure/FlinkProcedureITCase.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/procedure/FlinkProcedureITCase.java @@ -27,6 +27,7 @@ import org.apache.fluss.config.Configuration; import org.apache.fluss.config.MemorySize; import org.apache.fluss.exception.NoRebalanceInProgressException; +import org.apache.fluss.exception.RebalanceFailureException; import org.apache.fluss.exception.SecurityDisabledException; import org.apache.fluss.metadata.DataLakeFormat; import org.apache.fluss.metadata.TablePath; @@ -743,6 +744,40 @@ void testRebalance(boolean upperCase) throws Exception { FLUSS_CLUSTER_EXTENSION.getZooKeeperClient().deleteRebalanceTask(); } + @Test + void testPreferredLeaderElectionProcedure() throws Exception { + try { + String rebalance = + String.format( + "Call %s.sys.rebalance('PREFERRED_LEADER_ELECTION')", CATALOG_NAME); + try (CloseableIterator rows = tEnv.executeSql(rebalance).collect()) { + assertThat(CollectionUtil.iteratorToList(rows)).hasSize(1); + } + + retry( + Duration.ofMinutes(2), + () -> { + Optional progress = + admin.listRebalanceProgress(null).get(); + assertThat(progress).isPresent(); + assertThat(progress.get().status()).isEqualTo(RebalanceStatus.COMPLETED); + }); + + assertThatThrownBy( + () -> + tEnv.executeSql( + String.format( + "Call %s.sys.rebalance('PREFERRED_LEADER_ELECTION,LEADER_DISTRIBUTION')", + CATALOG_NAME)) + .await()) + .rootCause() + .isInstanceOf(RebalanceFailureException.class) + .hasMessageContaining("must be used as a standalone rebalance goal"); + } finally { + FLUSS_CLUSTER_EXTENSION.getZooKeeperClient().deleteRebalanceTask(); + } + } + @Test void testListRebalanceProgress() throws Exception { // add server tag PERMANENT_OFFLINE for server 3, this will avoid to generate bucket diff --git a/fluss-rust/crates/fluss/src/metadata/goal_type.rs b/fluss-rust/crates/fluss/src/metadata/goal_type.rs index 3880faae998..4a5379de341 100644 --- a/fluss-rust/crates/fluss/src/metadata/goal_type.rs +++ b/fluss-rust/crates/fluss/src/metadata/goal_type.rs @@ -23,6 +23,7 @@ pub enum GoalType { ReplicaDistribution, LeaderDistribution, RackAware, + PreferredLeaderElection, } impl GoalType { @@ -31,6 +32,7 @@ impl GoalType { Self::ReplicaDistribution => 0, Self::LeaderDistribution => 1, Self::RackAware => 2, + Self::PreferredLeaderElection => 3, } } @@ -39,6 +41,7 @@ impl GoalType { 0 => Ok(Self::ReplicaDistribution), 1 => Ok(Self::LeaderDistribution), 2 => Ok(Self::RackAware), + 3 => Ok(Self::PreferredLeaderElection), _ => Err(Error::IllegalArgument { message: format!("Unsupported GoalType: {value}"), }), @@ -56,6 +59,7 @@ mod tests { GoalType::ReplicaDistribution, GoalType::LeaderDistribution, GoalType::RackAware, + GoalType::PreferredLeaderElection, ] { assert_eq!(GoalType::try_from_i32(goal.to_i32()).unwrap(), goal); } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java index 0cff1592c4e..8935745467d 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java @@ -1605,6 +1605,23 @@ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) { coordinatorContext.getAssignment(tableBucket), newReplicas); if (planForBucket.isLeaderChanged() && !reassignment.isBeingReassigned()) { + int targetLeader = planForBucket.getNewLeader(); + Optional leaderAndIsr = + coordinatorContext.getBucketLeaderAndIsr(tableBucket); + Optional targetServerTag = coordinatorContext.getServerTag(targetLeader); + if (!leaderAndIsr.isPresent() + || !leaderAndIsr.get().isr().contains(targetLeader) + || !coordinatorContext.isReplicaOnline(targetLeader, tableBucket) + || (targetServerTag.isPresent() + && (targetServerTag.get() == ServerTag.TEMPORARY_OFFLINE + || targetServerTag.get() == ServerTag.PERMANENT_OFFLINE))) { + LOG.warn( + "Skipping leader-only rebalance for tableBucket {} because target leader {} is no longer eligible.", + tableBucket, + targetLeader); + rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED); + return; + } // buckets only need to change leader like leader replica rebalance. // Don't finish the task immediately; wait for the NotifyLeaderAndIsr response // from the tablet server to confirm the leader change has been applied. @@ -1614,7 +1631,7 @@ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) { tableBucketStateMachine.handleStateChange( Collections.singleton(tableBucket), OnlineBucket, - new ReassignmentLeaderElection(newReplicas)); + new ReassignmentLeaderElection(newReplicas, false)); } else { try { LOG.info( diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java index cc29b982b07..12bc2205ad0 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java @@ -346,7 +346,11 @@ public RebalanceTask generateRebalanceTask(List goalsByPriority) { try { // Generate the latest cluster model. long startTime = System.currentTimeMillis(); - ClusterModel clusterModel = buildClusterModel(eventProcessor.getCoordinatorContext()); + boolean preferredLeaderElection = + GoalOptimizer.isPreferredLeaderElection(goalsByPriority); + ClusterModel clusterModel = + buildClusterModel( + eventProcessor.getCoordinatorContext(), preferredLeaderElection); LOG.info( "Build cluster model for rebalance id {} with {} ms.", rebalanceId, @@ -420,8 +424,14 @@ private void completeRebalance() { } private ClusterModel buildClusterModel(CoordinatorContext coordinatorContext) { + return buildClusterModel(coordinatorContext, false); + } + + private ClusterModel buildClusterModel( + CoordinatorContext coordinatorContext, boolean preferredLeaderElection) { Map liveTabletServers = coordinatorContext.getLiveTabletServers(); Map serverTags = coordinatorContext.getServerTags(); + Set allBuckets = coordinatorContext.getAllBuckets(); Map serverModelMap = new HashMap<>(); for (ServerInfo serverInfo : liveTabletServers.values()) { @@ -434,11 +444,18 @@ private ClusterModel buildClusterModel(CoordinatorContext coordinatorContext) { serverModelMap.put(id, new ServerModel(id, rack, false)); } } + if (preferredLeaderElection) { + for (TableBucket tableBucket : allBuckets) { + for (Integer replica : coordinatorContext.getAssignment(tableBucket)) { + serverModelMap.putIfAbsent( + replica, new ServerModel(replica, RackModel.DEFAULT_RACK, true)); + } + } + } ClusterModel clusterModel = initialClusterModel(serverModelMap); // Try to update the cluster model with the latest bucket states. - Set allBuckets = coordinatorContext.getAllBuckets(); for (TableBucket tableBucket : allBuckets) { List assignment = coordinatorContext.getAssignment(tableBucket); Optional bucketLeaderAndIsrOpt = @@ -457,7 +474,17 @@ private ClusterModel buildClusterModel(CoordinatorContext coordinatorContext) { } for (int i = 0; i < assignment.size(); i++) { int replica = assignment.get(i); - clusterModel.createReplica(replica, tableBucket, i, leader == replica); + if (preferredLeaderElection) { + boolean isLeaderEligible = + liveTabletServers.containsKey(replica) + && coordinatorContext.isReplicaOnline(replica, tableBucket) + && isr.isr().contains(replica) + && !serverModelMap.get(replica).isOfflineTagged(); + clusterModel.createReplica( + replica, tableBucket, i, leader == replica, isLeaderEligible); + } else { + clusterModel.createReplica(replica, tableBucket, i, leader == replica); + } } } return clusterModel; diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/GoalOptimizer.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/GoalOptimizer.java index 6a34c0122ab..7f3493b93c6 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/GoalOptimizer.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/GoalOptimizer.java @@ -38,6 +38,7 @@ public class GoalOptimizer { public List doOptimizeOnce( ClusterModel clusterModel, List goalsByPriority) { + isPreferredLeaderElection(goalsByPriority); LOG.trace("Cluster before optimization is {}", clusterModel); Map> initReplicaDistribution = clusterModel.getReplicaDistribution(); @@ -77,4 +78,17 @@ public List doOptimizeOnce( return getDiff(initReplicaDistribution, initLeaderDistribution, clusterModel); } + + /** Validate the goal combination and return whether it is preferred leader election. */ + public static boolean isPreferredLeaderElection(List goalsByPriority) { + long preferredLeaderGoals = + goalsByPriority.stream() + .filter(PreferredLeaderElectionGoal.class::isInstance) + .count(); + if (preferredLeaderGoals > 0 && goalsByPriority.size() != 1) { + throw new IllegalArgumentException( + "PREFERRED_LEADER_ELECTION must be used as a standalone rebalance goal."); + } + return preferredLeaderGoals == 1; + } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/GoalUtils.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/GoalUtils.java index 652b2a2312f..f5f79540c1f 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/GoalUtils.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/GoalUtils.java @@ -39,6 +39,8 @@ public static Goal getGoalByType(GoalType goalType) { return new LeaderReplicaDistributionGoal(); case RACK_AWARE: return new RackAwareGoal(); + case PREFERRED_LEADER_ELECTION: + return new PreferredLeaderElectionGoal(); default: throw new IllegalArgumentException("Unsupported goal type " + goalType); } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/PreferredLeaderElectionGoal.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/PreferredLeaderElectionGoal.java new file mode 100644 index 00000000000..7f7afd076e6 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/goal/PreferredLeaderElectionGoal.java @@ -0,0 +1,95 @@ +/* + * 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.server.coordinator.rebalance.goal; + +import org.apache.fluss.server.coordinator.rebalance.ActionAcceptance; +import org.apache.fluss.server.coordinator.rebalance.ActionType; +import org.apache.fluss.server.coordinator.rebalance.RebalancingAction; +import org.apache.fluss.server.coordinator.rebalance.model.BucketModel; +import org.apache.fluss.server.coordinator.rebalance.model.ClusterModel; +import org.apache.fluss.server.coordinator.rebalance.model.ReplicaModel; + +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.apache.fluss.server.coordinator.rebalance.ActionAcceptance.ACCEPT; +import static org.apache.fluss.server.coordinator.rebalance.ActionAcceptance.REPLICA_REJECT; +import static org.apache.fluss.server.coordinator.rebalance.goal.GoalUtils.HardGoalStatsComparator; + +/** + * Goal to move each bucket's leadership to the first replica in its persisted assignment. + * + *

The goal never changes replica assignments and skips a bucket when its preferred replica is + * not currently eligible to become leader. + */ +public class PreferredLeaderElectionGoal implements Goal { + + @Override + public void optimize(ClusterModel clusterModel, Set optimizedGoals) { + for (Map.Entry> entry : + clusterModel.getBucketsByTable().entrySet()) { + for (BucketModel bucket : entry.getValue()) { + List replicas = bucket.replicas(); + ReplicaModel leader = bucket.leader(); + if (replicas.isEmpty() || leader == null) { + continue; + } + + ReplicaModel preferredLeader = replicas.get(0); + if (leader.equals(preferredLeader) + || !preferredLeader.isLeaderEligible() + || preferredLeader.server().isOfflineTagged()) { + continue; + } + + clusterModel.relocateLeadership( + bucket.tableBucket(), leader.serverId(), preferredLeader.serverId()); + } + } + } + + @Override + public ActionAcceptance actionAcceptance(RebalancingAction action, ClusterModel clusterModel) { + if (action.getActionType() != ActionType.LEADERSHIP_MOVEMENT) { + return REPLICA_REJECT; + } + + BucketModel bucket = clusterModel.bucket(action.getTableBucket()); + if (bucket == null || bucket.replicas().isEmpty()) { + return REPLICA_REJECT; + } + + return bucket.replicas().get(0).serverId() == action.getDestinationServerId() + ? ACCEPT + : REPLICA_REJECT; + } + + @Override + public ClusterModelStatsComparator clusterModelStatsComparator() { + return new HardGoalStatsComparator(); + } + + @Override + public void finish() {} + + @Override + public String name() { + return getClass().getSimpleName(); + } +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/model/ClusterModel.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/model/ClusterModel.java index 8d481ecfd7f..60b8f92c3d5 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/model/ClusterModel.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/model/ClusterModel.java @@ -210,12 +210,21 @@ public Map getLeaderDistribution() { } public void createReplica(int serverId, TableBucket tableBucket, int index, boolean isLeader) { + createReplica(serverId, tableBucket, index, isLeader, true); + } + + public void createReplica( + int serverId, + TableBucket tableBucket, + int index, + boolean isLeader, + boolean isLeaderEligible) { ServerModel server = server(serverId); if (server == null) { throw new IllegalArgumentException("Server is not in the cluster."); } - ReplicaModel replica = new ReplicaModel(tableBucket, server, isLeader); + ReplicaModel replica = new ReplicaModel(tableBucket, server, isLeader, isLeaderEligible); server.putReplica(tableBucket, replica); if (!bucketsByTableBucket.containsKey(tableBucket)) { diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/model/ReplicaModel.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/model/ReplicaModel.java index e67d9bd733f..f39f8935510 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/model/ReplicaModel.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/model/ReplicaModel.java @@ -26,14 +26,24 @@ public class ReplicaModel { private final TableBucket tableBucket; private final ServerModel originalServer; + private final boolean isLeaderEligible; private ServerModel server; private boolean isLeader; public ReplicaModel(TableBucket tableBucket, ServerModel server, boolean isLeader) { + this(tableBucket, server, isLeader, true); + } + + public ReplicaModel( + TableBucket tableBucket, + ServerModel server, + boolean isLeader, + boolean isLeaderEligible) { this.tableBucket = tableBucket; this.server = server; this.isLeader = isLeader; this.originalServer = server; + this.isLeaderEligible = isLeaderEligible; } public TableBucket tableBucket() { @@ -56,6 +66,10 @@ public boolean isLeader() { return isLeader; } + public boolean isLeaderEligible() { + return isLeaderEligible; + } + public void makeFollower() { setLeadership(false); } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/statemachine/ReplicaLeaderElection.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/statemachine/ReplicaLeaderElection.java index d2e01eb8cc0..f27e0281653 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/statemachine/ReplicaLeaderElection.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/statemachine/ReplicaLeaderElection.java @@ -115,9 +115,15 @@ public Optional leaderElection( /** The reassignment replica leader election. */ public static class ReassignmentLeaderElection extends ReplicaLeaderElection { private final List newReplicas; + private final boolean allowFallback; public ReassignmentLeaderElection(List newReplicas) { + this(newReplicas, true); + } + + public ReassignmentLeaderElection(List newReplicas, boolean allowFallback) { this.newReplicas = newReplicas; + this.allowFallback = allowFallback; } public Optional leaderElection( @@ -139,6 +145,9 @@ public Optional leaderElection( if (availableReplicas.isEmpty()) { return Optional.empty(); } + if (!allowFallback && !availableReplicas.get(0).equals(newReplicas.get(0))) { + return Optional.empty(); + } // Always use the first available replica as leader to honor the rebalance plan. int newLeader = availableReplicas.get(0); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/goal/PreferredLeaderElectionGoalTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/goal/PreferredLeaderElectionGoalTest.java new file mode 100644 index 00000000000..62e42b318f6 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/goal/PreferredLeaderElectionGoalTest.java @@ -0,0 +1,138 @@ +/* + * 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.server.coordinator.rebalance.goal; + +import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.server.coordinator.rebalance.model.ClusterModel; +import org.apache.fluss.server.coordinator.rebalance.model.ServerModel; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.SortedSet; +import java.util.TreeSet; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Test for {@link PreferredLeaderElectionGoal}. */ +class PreferredLeaderElectionGoalTest { + + private static final TableBucket TABLE_BUCKET = new TableBucket(1L, 0); + private static final List ASSIGNMENT = Arrays.asList(0, 1, 2); + + @Test + void testGenerateLeaderOnlyPlanAndRemainIdempotent() { + ClusterModel clusterModel = createClusterModel(false); + addBucket(clusterModel, TABLE_BUCKET, 1, true); + + GoalOptimizer optimizer = new GoalOptimizer(); + List plans = + optimizer.doOptimizeOnce( + clusterModel, Collections.singletonList(new PreferredLeaderElectionGoal())); + + assertThat(plans) + .containsExactly( + new RebalancePlanForBucket(TABLE_BUCKET, 1, 0, ASSIGNMENT, ASSIGNMENT)); + assertThat(clusterModel.getReplicaDistribution().get(TABLE_BUCKET)) + .containsExactlyElementsOf(ASSIGNMENT); + assertThat(clusterModel.getLeaderDistribution().get(TABLE_BUCKET)).isEqualTo(0); + + assertThat( + optimizer.doOptimizeOnce( + clusterModel, + Collections.singletonList(new PreferredLeaderElectionGoal()))) + .isEmpty(); + } + + @Test + void testSkipIneligiblePreferredLeaderWithoutFallback() { + ClusterModel clusterModel = createClusterModel(false); + addBucket(clusterModel, TABLE_BUCKET, 1, false); + + List plans = + new GoalOptimizer() + .doOptimizeOnce( + clusterModel, + Collections.singletonList(new PreferredLeaderElectionGoal())); + + assertThat(plans).isEmpty(); + assertThat(clusterModel.getLeaderDistribution().get(TABLE_BUCKET)).isEqualTo(1); + assertThat(clusterModel.getReplicaDistribution().get(TABLE_BUCKET)) + .containsExactlyElementsOf(ASSIGNMENT); + } + + @Test + void testSkipOfflineTaggedPreferredLeader() { + ClusterModel clusterModel = createClusterModel(true); + addBucket(clusterModel, TABLE_BUCKET, 1, true); + + List plans = + new GoalOptimizer() + .doOptimizeOnce( + clusterModel, + Collections.singletonList(new PreferredLeaderElectionGoal())); + + assertThat(plans).isEmpty(); + assertThat(clusterModel.getLeaderDistribution().get(TABLE_BUCKET)).isEqualTo(1); + } + + @Test + void testRejectGoalCombination() { + ClusterModel clusterModel = createClusterModel(false); + + assertThatThrownBy( + () -> + new GoalOptimizer() + .doOptimizeOnce( + clusterModel, + Arrays.asList( + new PreferredLeaderElectionGoal(), + new LeaderReplicaDistributionGoal()))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage( + "PREFERRED_LEADER_ELECTION must be used as a standalone rebalance goal."); + } + + private static ClusterModel createClusterModel(boolean preferredLeaderOfflineTagged) { + SortedSet servers = new TreeSet<>(); + servers.add(new ServerModel(0, "rack0", preferredLeaderOfflineTagged)); + servers.add(new ServerModel(1, "rack1", false)); + servers.add(new ServerModel(2, "rack2", false)); + return new ClusterModel(servers); + } + + private static void addBucket( + ClusterModel clusterModel, + TableBucket tableBucket, + int currentLeader, + boolean preferredLeaderEligible) { + for (int i = 0; i < ASSIGNMENT.size(); i++) { + int replica = ASSIGNMENT.get(i); + clusterModel.createReplica( + replica, + tableBucket, + i, + replica == currentLeader, + replica != ASSIGNMENT.get(0) || preferredLeaderEligible); + } + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/statemachine/ReplicaLeaderElectionTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/statemachine/ReplicaLeaderElectionTest.java index 313114d69d6..5db0c8a9d5e 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/statemachine/ReplicaLeaderElectionTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/statemachine/ReplicaLeaderElectionTest.java @@ -228,6 +228,10 @@ void testReassignBucketLeaderElection() { assertThat(leaderOpt).isPresent(); assertThat(leaderOpt.get().getLeaderAndIsr().leader()).isEqualTo(2); + reassignmentLeaderElection = new ReassignmentLeaderElection(targetReplicas, false); + leaderOpt = reassignmentLeaderElection.leaderElection(liveReplicas, leaderAndIsr, false); + assertThat(leaderOpt).isEmpty(); + targetReplicas = Arrays.asList(1, 2, 3); reassignmentLeaderElection = new ReassignmentLeaderElection(targetReplicas); liveReplicas = Arrays.asList(1, 2); diff --git a/website/docs/maintenance/operations/rebalance.md b/website/docs/maintenance/operations/rebalance.md index 418b7c585f2..701efb0ec33 100644 --- a/website/docs/maintenance/operations/rebalance.md +++ b/website/docs/maintenance/operations/rebalance.md @@ -72,6 +72,23 @@ Available rebalance goals: - **RACK_AWARE**: Ensures replicas of the same bucket are distributed across different racks. This goal is essential for high availability in multi-rack deployments, as it prevents data loss when an entire rack fails. - **REPLICA_DISTRIBUTION**: Ensures the number of replicas on each TabletServer is near balanced - **LEADER_DISTRIBUTION**: Ensures the number of leader replicas on each TabletServer is near balanced +- **PREFERRED_LEADER_ELECTION**: Moves each bucket's leadership to the first replica in its persisted assignment without changing the replica set. Buckets are skipped when the preferred replica is unavailable, offline, outside the ISR, or tagged `TEMPORARY_OFFLINE` or `PERMANENT_OFFLINE`. + +`PREFERRED_LEADER_ELECTION` must be requested by itself: + +```java +String rebalanceId = admin.rebalance( + Collections.singletonList(GoalType.PREFERRED_LEADER_ELECTION) +).get(); +``` + +The same operation is available through Flink SQL: + +```sql +CALL sys.rebalance('PREFERRED_LEADER_ELECTION'); +``` + +Combining this goal with another rebalance goal is rejected. :::tip Goal Priority Goals are processed in the order specified. When using `RACK_AWARE`, always place it first to ensure subsequent goals (like `REPLICA_DISTRIBUTION`) respect rack constraints. If `RACK_AWARE` is not the first goal, replica movements may violate rack awareness requirements.