KAFKA-19048: Minimal Movement Replica Balancing algorithm

The new replica rebalancing strategy aims to achieve the following objectives:

1. Minimal Movement: Minimize the number of replica relocations during rebalancing.
2. Replica Balancing: Ensure that replicas are evenly distributed across brokers.
3. Anti-Affinity Support: Support rack-aware allocation when enabled.
4. Leader Balancing: Distribute leader replicas evenly across brokers.
5. ISR Order Optimization: Optimize adjacency relationships to prevent failover traffic concentration in case of broker failures.

Author: Jialun Peng

metadata/src/main/java/org/apache/kafka/metadata/MinimalMovementReplicaBalancer.java

@@ -1,11 +1,28 @@
+/*
+ * 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.kafka.metadata;
 
-import org.apache.kafka.admin.BrokerMetadata;
 import org.apache.kafka.common.errors.BrokerNotAvailableException;
 import org.apache.kafka.common.errors.DuplicateBrokerRegistrationException;
 import org.apache.kafka.common.errors.InvalidPartitionsException;
 import org.apache.kafka.common.errors.InvalidReplicaAssignmentException;
 import org.apache.kafka.common.errors.InvalidReplicationFactorException;
+import org.apache.kafka.metadata.placement.UsableBroker;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -60,7 +77,7 @@
  *
  *   <li><b>Node-Level Replica Distribution:</b>
  *     <ul>
-        <p>If the number of replicas assigned to a rack is not a multiple of the number of nodes in that rack, some nodes will host one additional replica compared to others.</p>
+ * <p>If the number of replicas assigned to a rack is not a multiple of the number of nodes in that rack, some nodes will host one additional replica compared to others.</p>
  *       <li><b>Rack Count = Replication Factor:</b>
  *         <ul>
  *           <li>If all racks contain an equal number of nodes, each node will have the same number of replicas.</li>
@@ -113,7 +130,7 @@ public class MinimalMovementReplicaBalancer {
      */
     private final int replicationFactor;
 
-    private final List<BrokerMetadata> brokerMetadatas;
+    private final List<UsableBroker> usableBrokers;
 
     /**
      * Total number of replicas across all partitions (calculated as partitionCount * replicationFactor).
@@ -177,7 +194,7 @@ public class MinimalMovementReplicaBalancer {
      */
     private final Map<String, Integer> rackToRemainderReplicaCount = new HashMap<>();
 
-    private final static String COMMON_REMAINDER_RACK = UUID.randomUUID().toString();
+    private static final String COMMON_REMAINDER_RACK = UUID.randomUUID().toString();
 
     /**
      * <p>Partition assignments per broker, tracking partition locations.</p>
@@ -205,20 +222,20 @@ public class MinimalMovementReplicaBalancer {
 
     public static final String BROKER_ADJACENCY_SYMBOL = "->";
 
-    public MinimalMovementReplicaBalancer(Map<Integer, List<Integer>> assignment, List<Integer> targetBrokerIds, List<BrokerMetadata> brokerMetadatas, boolean enableRackAwareness) {
+    public MinimalMovementReplicaBalancer(Map<Integer, List<Integer>> assignment, List<Integer> targetBrokerIds, List<UsableBroker> usableBrokers, boolean enableRackAwareness) {
         if (assignment == null || assignment.isEmpty()) {
             throw new InvalidReplicaAssignmentException("assignment is empty.");
         }
         if (targetBrokerIds == null) {
             throw new BrokerNotAvailableException("targetBrokerIds is null.");
         }
         // Nodes with unclear rack information in the proxy metadata are not considered as anti-affinity.
-        if (brokerMetadatas == null) {
-            this.brokerMetadatas = new ArrayList<>();
+        if (usableBrokers == null) {
+            this.usableBrokers = new ArrayList<>();
         } else {
-            this.brokerMetadatas = new ArrayList<>(brokerMetadatas);
+            this.usableBrokers = new ArrayList<>(usableBrokers);
         }
-        if (this.brokerMetadatas.stream().map(broker -> broker.id).collect(Collectors.toSet()).size() < this.brokerMetadatas.size()) {
+        if (this.usableBrokers.stream().map(UsableBroker::id).collect(Collectors.toSet()).size() < this.usableBrokers.size()) {
             throw new DuplicateBrokerRegistrationException("Duplicate broker ID found.");
         }
         this.assignment = assignment.entrySet().stream()
@@ -245,7 +262,7 @@ private void verifyingParameters() {
             .distinct()
             .count();
         if (replicationFactor > rackCount)
-            throw new InvalidReplicationFactorException("Replication factor: " + replicationFactor + " larger than available brokers or rack: " + rackCount + ".");
+            throw new InvalidReplicationFactorException("Replication factor: " + replicationFactor + " larger than available brokers or rack num: " + rackCount + ".");
     }
 
     private void initializeClusterState() {
@@ -567,8 +584,8 @@ private boolean hasReplicaMoved(int brokerId, List<Integer> replicas, boolean ha
      */
     private void balanceLeadersAcrossBrokers() {
         Map<Integer, Integer> brokerLeaderCount = new HashMap<>();
-        for (Integer partition : assignment.keySet()) {
-            List<Integer> replicas = assignment.get(partition);
+        for (Map.Entry<Integer, List<Integer>> entry : assignment.entrySet()) {
+            List<Integer> replicas = entry.getValue();
             int leaderId = replicas.get(0);
             Integer currentLeaderCount = brokerLeaderCount.getOrDefault(leaderId, 1);
             Integer idealLeaderCount = brokerIdealLeaderCount.get(leaderId);
@@ -772,16 +789,16 @@ private void identifyBrokerPartitionAssignments() {
      * This unifies the logic for both rack-aware and non-rack-aware scenarios.
      */
     private void computeBrokerRackMap() {
-        Set<Integer> existingBrokerIds = brokerMetadatas.stream()
-            .map(b -> b.id)
+        Set<Integer> existingBrokerIds = usableBrokers.stream()
+            .map(UsableBroker::id)
             .collect(Collectors.toSet());
-        involvedBrokerIds.stream().filter(brokerId -> !existingBrokerIds.contains(brokerId)).forEach(unknownBrokerId -> brokerMetadatas.add(new BrokerMetadata(unknownBrokerId, Optional.of(UUID.randomUUID().toString()))));
-        brokerRackMapping = brokerMetadatas.stream()
-            .filter(b -> involvedBrokerIds.contains(b.id))
+        involvedBrokerIds.stream().filter(brokerId -> !existingBrokerIds.contains(brokerId)).forEach(unknownBrokerId -> usableBrokers.add(new UsableBroker(unknownBrokerId, Optional.of(UUID.randomUUID().toString()), false)));
+        brokerRackMapping = usableBrokers.stream()
+            .filter(b -> involvedBrokerIds.contains(b.id()))
             .collect(Collectors.toMap(
-                broker -> broker.id,
+                UsableBroker::id,
                 broker -> enableRackAwareness
-                    ? broker.rack.orElse(UUID.randomUUID().toString())
+                    ? broker.rack().orElse(UUID.randomUUID().toString())
                     : UUID.randomUUID().toString()
             ));
     }