Port TopicLeaderDistributionGoal and ExecutionTaskPlanner changes

cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/analyzer/BalancingConstraint.java

@@ -28,6 +28,7 @@ public class BalancingConstraint {
   private final double _topicReplicaBalancePercentage;
   private final int _topicReplicaBalanceMinGap;
   private final int _topicReplicaBalanceMaxGap;
+  private final double _topicLeaderBalancePercentage;
   private final double _goalViolationDistributionThresholdMultiplier;
   private final Map<Resource, Double> _capacityThreshold;
   private final Map<Resource, Double> _lowUtilizationThreshold;
@@ -76,6 +77,7 @@ public BalancingConstraint(KafkaCruiseControlConfig config) {
     _replicaBalancePercentage = config.getDouble(AnalyzerConfig.REPLICA_COUNT_BALANCE_THRESHOLD_CONFIG);
     _leaderReplicaBalancePercentage = config.getDouble(AnalyzerConfig.LEADER_REPLICA_COUNT_BALANCE_THRESHOLD_CONFIG);
     _topicReplicaBalancePercentage = config.getDouble(AnalyzerConfig.TOPIC_REPLICA_COUNT_BALANCE_THRESHOLD_CONFIG);
+    _topicLeaderBalancePercentage = config.getDouble(AnalyzerConfig.TOPIC_LEADER_COUNT_BALANCE_THRESHOLD_CONFIG);
     _topicReplicaBalanceMinGap = config.getInt(AnalyzerConfig.TOPIC_REPLICA_COUNT_BALANCE_MIN_GAP_CONFIG);
     _topicReplicaBalanceMaxGap = config.getInt(AnalyzerConfig.TOPIC_REPLICA_COUNT_BALANCE_MAX_GAP_CONFIG);
     _goalViolationDistributionThresholdMultiplier = config.getDouble(AnalyzerConfig.GOAL_VIOLATION_DISTRIBUTION_THRESHOLD_MULTIPLIER_CONFIG);
@@ -122,6 +124,7 @@ Properties setProps(Properties props) {
     props.put(AnalyzerConfig.TOPIC_REPLICA_COUNT_BALANCE_THRESHOLD_CONFIG, Double.toString(_topicReplicaBalancePercentage));
     props.put(AnalyzerConfig.TOPIC_REPLICA_COUNT_BALANCE_MIN_GAP_CONFIG, Integer.toString(_topicReplicaBalanceMinGap));
     props.put(AnalyzerConfig.TOPIC_REPLICA_COUNT_BALANCE_MAX_GAP_CONFIG, Integer.toString(_topicReplicaBalanceMaxGap));
+    props.put(AnalyzerConfig.TOPIC_LEADER_COUNT_BALANCE_THRESHOLD_CONFIG, Double.toString(_topicLeaderBalancePercentage));
     props.put(AnalyzerConfig.GOAL_VIOLATION_DISTRIBUTION_THRESHOLD_MULTIPLIER_CONFIG, Double.toString(_goalViolationDistributionThresholdMultiplier));
     props.put(AnalyzerConfig.TOPICS_WITH_MIN_LEADERS_PER_BROKER_CONFIG, _topicsWithMinLeadersPerBrokerPattern.pattern());
     props.put(AnalyzerConfig.MIN_TOPIC_LEADERS_PER_BROKER_CONFIG, Integer.toString(_minTopicLeadersPerBroker));
@@ -197,6 +200,13 @@ public int topicReplicaBalanceMaxGap() {
     return _topicReplicaBalanceMaxGap;
   }
 
+  /**
+   * @return Topic leader replica balance percentage for {@link com.linkedin.kafka.cruisecontrol.analyzer.goals.TopicLeaderDistributionGoal}.
+   */
+  public double topicLeaderBalancePercentage() {
+    return _topicLeaderBalancePercentage;
+  }
+
   /**
    * @return Goal violation distribution threshold multiplier to be used in detection and fixing goal violations.
    */

cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/config/constants/AnalyzerConfig.java

@@ -133,6 +133,14 @@ public final class AnalyzerConfig {
       + " the average replica count for each topic. A balance limit is set via topic.replica.count.balance.threshold config."
       + " If the difference between the computed limit and the average replica count for the relevant topic is greater than"
       + " the value specified by this config, the limit is adjusted accordingly.";
+  /**
+   * <code>topic.leader.count.balance.threshold</code>
+   */
+  public static final String TOPIC_LEADER_COUNT_BALANCE_THRESHOLD_CONFIG = "topic.leader.count.balance.threshold";
+  public static final double DEFAULT_TOPIC_LEADER_COUNT_BALANCE_THRESHOLD = 3.0;
+  public static final String TOPIC_LEADER_COUNT_BALANCE_THRESHOLD_DOC = "The maximum allowed extent of unbalance for "
+      + "leader replica distribution from each topic. For example, 1.80 means the highest topic leader replica count of a "
+      + "broker should not be above 1.80x of average leader replica count of all brokers for the same topic.";
 
   /**
    * <code>cpu.capacity.threshold</code>
@@ -524,6 +532,12 @@ public static ConfigDef define(ConfigDef configDef) {
                             atLeast(1),
                             ConfigDef.Importance.MEDIUM,
                             TOPIC_REPLICA_COUNT_BALANCE_MAX_GAP_DOC)
+                    .define(TOPIC_LEADER_COUNT_BALANCE_THRESHOLD_CONFIG,
+                            ConfigDef.Type.DOUBLE,
+                            DEFAULT_TOPIC_LEADER_COUNT_BALANCE_THRESHOLD,
+                            atLeast(1),
+                            ConfigDef.Importance.HIGH,
+                            TOPIC_LEADER_COUNT_BALANCE_THRESHOLD_DOC)
                     .define(CPU_CAPACITY_THRESHOLD_CONFIG,
                             ConfigDef.Type.DOUBLE,
                             DEFAULT_CPU_CAPACITY_THRESHOLD,

cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlanner.java

@@ -5,7 +5,6 @@
 package com.linkedin.kafka.cruisecontrol.executor;
 
 import com.linkedin.cruisecontrol.common.utils.Utils;
-import com.linkedin.kafka.cruisecontrol.KafkaCruiseControlUtils;
 import com.linkedin.kafka.cruisecontrol.config.KafkaCruiseControlConfig;
 import com.linkedin.kafka.cruisecontrol.executor.concurrency.ExecutionConcurrencyManager;
 import com.linkedin.kafka.cruisecontrol.executor.strategy.BaseReplicaMovementStrategy;
@@ -398,8 +397,7 @@ public List<ExecutionTask> getInterBrokerReplicaMovementTasks(Map<Integer, Integ
           int sourceBroker = task.proposal().oldLeader().brokerId();
           Set<Integer> destinationBrokers = task.proposal().replicasToAdd().stream().mapToInt(ReplicaPlacementInfo::brokerId)
                                                 .boxed().collect(Collectors.toSet());
-          if (brokerInvolved.contains(sourceBroker)
-              || KafkaCruiseControlUtils.containsAny(brokerInvolved, destinationBrokers)) {
+          if (brokerInvolved.contains(sourceBroker)) {
             continue;
           }
           TopicPartition tp = task.proposal().topicPartition();

cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/model/Broker.java

@@ -23,6 +23,7 @@
 import java.util.SortedMap;
 import java.util.TreeMap;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 import org.apache.kafka.common.TopicPartition;
 
 import static com.linkedin.cruisecontrol.common.utils.Utils.validateNotNull;
@@ -203,6 +204,15 @@ public int numLeadersFor(String topicName) {
     return (int) replicasOfTopicInBroker(topicName).stream().filter(Replica::isLeader).count();
   }
 
+  /**
+   * Get leader replicas for topic.
+   * @param topic Topic of the requested replicas.
+   * @return Leader replicas in this broker sharing the given topic.
+   */
+  public Collection<Replica> leadersFor(String topic) {
+    return replicasOfTopicInBroker(topic).stream().filter(Replica::isLeader).collect(Collectors.toList());
+  }
+
   /**
    * @return {@code true} if the broker is not dead, {@code false} otherwise.
    */

cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/model/ClusterModel.java

@@ -671,6 +671,19 @@ public int numTopicReplicas(String topic) {
     return _numReplicasByTopic.getOrDefault(topic, 0);
   }
 
+  /**
+   * Get the number of leader replicas with the given topic name in cluster.
+   * @param topic Name of the topic for which the number of leader replicas in cluster will be counted.
+   * @return Number of leader replicas with the given topic name in cluster.
+   */
+  public int numTopicLeaders(String topic) {
+    int numTopicLeaders = 0;
+    for (Broker broker : brokers()) {
+      numTopicLeaders += broker.numLeadersFor(topic);
+    }
+    return numTopicLeaders;
+  }
+
   /**
    * Get the number of leader replicas in cluster.
    *

cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/model/ClusterModelStats.java

@@ -37,6 +37,7 @@ public class ClusterModelStats {
   private final Map<Statistic, Number> _replicaStats;
   private final Map<Statistic, Number> _leaderReplicaStats;
   private final Map<Statistic, Number> _topicReplicaStats;
+  private final Map<Statistic, Number> _topicLeaderStats;
   private int _numBrokers;
   private int _numReplicasInCluster;
   private int _numPartitionsWithOfflineReplicas;
@@ -62,6 +63,7 @@ public class ClusterModelStats {
     _replicaStats = new HashMap<>();
     _leaderReplicaStats = new HashMap<>();
     _topicReplicaStats = new HashMap<>();
+    _topicLeaderStats = new HashMap<>();
     _numBrokers = 0;
     _numReplicasInCluster = 0;
     _numPartitionsWithOfflineReplicas = 0;
@@ -94,6 +96,7 @@ ClusterModelStats populate(ClusterModel clusterModel, BalancingConstraint balanc
     numForReplicas(clusterModel, brokers, aliveBrokers);
     numForLeaderReplicas(brokers, aliveBrokers);
     numForAvgTopicReplicas(clusterModel, brokers, topics);
+    numForAvgTopicLeaders(clusterModel, brokers, topics, aliveBrokers);
     _utilizationMatrix = clusterModel.utilizationMatrix();
     _numSnapshotWindows = clusterModel.load().numWindows();
     _monitoredPartitionsRatio = clusterModel.monitoredPartitionsRatio();
@@ -136,6 +139,13 @@ public Map<Statistic, Number> topicReplicaStats() {
     return Collections.unmodifiableMap(_topicReplicaStats);
   }
 
+  /**
+   * @return Topic leader stats for the cluster instance that the object was populated with.
+   */
+  public Map<Statistic, Number> topicLeaderStats() {
+    return Collections.unmodifiableMap(_topicLeaderStats);
+  }
+
   /**
    * @return The number of brokers for the cluster instance that the object was populated with.
    */
@@ -475,6 +485,49 @@ private void numForAvgTopicReplicas(ClusterModel clusterModel, SortedSet<Broker>
     _topicReplicaStats.put(Statistic.ST_DEV, _topicReplicaStats.get(Statistic.ST_DEV).doubleValue() / _numTopics);
   }
 
+  /**
+   * Generate statistics for leader replicas of each topic in the given cluster.
+   * Average and standard deviation calculations are based on alive brokers.
+   * @param clusterModel The state of the cluster.
+   * @param brokers Brokers in the cluster.
+   * @param topics Topics in the cluster.
+   * @param aliveBrokers Alive brokers in the cluster.
+   */
+  private void numForAvgTopicLeaders(ClusterModel clusterModel,
+                                     SortedSet<Broker> brokers,
+                                     Set<String> topics,
+                                     Set<Broker> aliveBrokers) {
+    _topicLeaderStats.put(Statistic.AVG, 0.0);
+    _topicLeaderStats.put(Statistic.MAX, 0);
+    _topicLeaderStats.put(Statistic.MIN, Integer.MAX_VALUE);
+    _topicLeaderStats.put(Statistic.ST_DEV, 0.0);
+    int numAliveBrokers = aliveBrokers.size();
+    for (String topic : topics) {
+      int maxTopicLeadersInBroker = 0;
+      int minTopicLeadersInBroker = Integer.MAX_VALUE;
+      double avgTopicLeaders = ((double) clusterModel.numTopicLeaders(topic)) / numAliveBrokers;
+      double variance = 0.0;
+      for (Broker broker : brokers) {
+        int numTopicLeadersInBroker = broker.numLeadersFor(topic);
+        maxTopicLeadersInBroker = Math.max(maxTopicLeadersInBroker, numTopicLeadersInBroker);
+        minTopicLeadersInBroker = Math.min(minTopicLeadersInBroker, numTopicLeadersInBroker);
+        if (broker.isAlive()) {
+          // Standard deviation of leader replicas in alive brokers.
+          variance += (Math.pow(numTopicLeadersInBroker - avgTopicLeaders, 2) / numAliveBrokers);
+        }
+      }
+      _topicLeaderStats.put(Statistic.AVG, _topicLeaderStats.get(Statistic.AVG).doubleValue() + avgTopicLeaders);
+      _topicLeaderStats.put(Statistic.MAX,
+                            Math.max(_topicLeaderStats.get(Statistic.MAX).intValue(), maxTopicLeadersInBroker));
+      _topicLeaderStats.put(Statistic.MIN,
+                            Math.min(_topicLeaderStats.get(Statistic.MIN).intValue(), minTopicLeadersInBroker));
+      _topicLeaderStats.put(Statistic.ST_DEV, (Double) _topicLeaderStats.get(Statistic.ST_DEV) + Math.sqrt(variance));
+    }
+
+    _topicLeaderStats.put(Statistic.AVG, _topicLeaderStats.get(Statistic.AVG).doubleValue() / _numTopics);
+    _topicLeaderStats.put(Statistic.ST_DEV, _topicLeaderStats.get(Statistic.ST_DEV).doubleValue() / _numTopics);
+  }
+
   /**
    * Generate statistics for disks in the given cluster.
    * For each alive disk on disk broker in the cluster, check whether its utilization percentage is within the range centered

cruise-control/src/test/java/com/linkedin/kafka/cruisecontrol/analyzer/DeterministicClusterTest.java

@@ -24,6 +24,7 @@
 import com.linkedin.kafka.cruisecontrol.analyzer.goals.RackAwareGoal;
 import com.linkedin.kafka.cruisecontrol.analyzer.goals.ReplicaCapacityGoal;
 import com.linkedin.kafka.cruisecontrol.analyzer.goals.ReplicaDistributionGoal;
+import com.linkedin.kafka.cruisecontrol.analyzer.goals.TopicLeaderDistributionGoal;
 import com.linkedin.kafka.cruisecontrol.analyzer.goals.TopicReplicaDistributionGoal;
 import com.linkedin.kafka.cruisecontrol.analyzer.kafkaassigner.KafkaAssignerDiskUsageDistributionGoal;
 import com.linkedin.kafka.cruisecontrol.analyzer.kafkaassigner.KafkaAssignerEvenRackAwareGoal;
@@ -115,6 +116,7 @@ public static Collection<Object[]> data() {
                                                     LeaderReplicaDistributionGoal.class.getName(),
                                                     LeaderBytesInDistributionGoal.class.getName(),
                                                     TopicReplicaDistributionGoal.class.getName(),
+                                                    TopicLeaderDistributionGoal.class.getName(),
                                                     PreferredLeaderElectionGoal.class.getName());
 
     List<OptimizationVerifier.Verification> verifications = Arrays.asList(NEW_BROKERS, BROKEN_BROKERS, REGRESSION);

cruise-control/src/test/java/com/linkedin/kafka/cruisecontrol/analyzer/RandomClusterTest.java

@@ -25,6 +25,7 @@
 import com.linkedin.kafka.cruisecontrol.analyzer.goals.ReplicaCapacityGoal;
 import com.linkedin.kafka.cruisecontrol.analyzer.goals.ReplicaDistributionGoal;
 import com.linkedin.kafka.cruisecontrol.analyzer.goals.TopicReplicaDistributionGoal;
+import com.linkedin.kafka.cruisecontrol.analyzer.goals.TopicLeaderDistributionGoal;
 import com.linkedin.kafka.cruisecontrol.analyzer.kafkaassigner.KafkaAssignerDiskUsageDistributionGoal;
 import com.linkedin.kafka.cruisecontrol.analyzer.kafkaassigner.KafkaAssignerEvenRackAwareGoal;
 import com.linkedin.kafka.cruisecontrol.common.Resource;
@@ -118,6 +119,7 @@ public static Collection<Object[]> data(TestConstants.Distribution distribution)
                                                     LeaderReplicaDistributionGoal.class.getName(),
                                                     LeaderBytesInDistributionGoal.class.getName(),
                                                     TopicReplicaDistributionGoal.class.getName(),
+                                                    TopicLeaderDistributionGoal.class.getName(),
                                                     PreferredLeaderElectionGoal.class.getName());
 
     List<String> kafkaAssignerGoals = Arrays.asList(KafkaAssignerEvenRackAwareGoal.class.getName(),

cruise-control/src/test/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlannerTest.java

@@ -288,8 +288,8 @@ public void testGetInterBrokerPartitionMovementTasks() {
     partitionMovementTasks = prioritizeSmallMovementPlanner.getInterBrokerReplicaMovementTasks(readyBrokers, Collections.emptySet(),
                                                                                                _defaultPartitionsMaxCap);
     assertEquals("First task", _partitionMovement0, partitionMovementTasks.get(0).proposal());
-    assertEquals("Second task", _partitionMovement2, partitionMovementTasks.get(1).proposal());
-    assertEquals("Third task", _partitionMovement3, partitionMovementTasks.get(2).proposal());
+    assertEquals("Second task", _partitionMovement3, partitionMovementTasks.get(1).proposal());
+    assertEquals("Third task", _partitionMovement2, partitionMovementTasks.get(2).proposal());
     assertEquals("Fourth task", _partitionMovement1, partitionMovementTasks.get(3).proposal());
 
     smallUrpMovementPlanner.addExecutionProposals(proposals, strategyOptions, null);
@@ -356,8 +356,8 @@ public void testGetInterBrokerPartitionMovementWithMinIsrTasks() {
     List<ExecutionTask> partitionMovementTasks
         = prioritizeOneAboveMinIsrMovementPlanner.getInterBrokerReplicaMovementTasks(readyBrokers, Collections.emptySet(), _defaultPartitionsMaxCap);
     assertEquals("First task", _rf4PartitionMovement2, partitionMovementTasks.get(0).proposal());
-    assertEquals("Second task", _rf4PartitionMovement3, partitionMovementTasks.get(1).proposal());
-    assertEquals("Third task", _rf4PartitionMovement1, partitionMovementTasks.get(2).proposal());
+    assertEquals("Second task", _rf4PartitionMovement1, partitionMovementTasks.get(1).proposal());
+    assertEquals("Third task", _rf4PartitionMovement3, partitionMovementTasks.get(2).proposal());
     assertEquals("Fourth task", _rf4PartitionMovement0, partitionMovementTasks.get(3).proposal());
   }
 
@@ -407,8 +407,8 @@ public void testDynamicConfigReplicaMovementStrategy() {
     planner.addExecutionProposals(proposals, strategyOptions, new PrioritizeSmallReplicaMovementStrategy());
     partitionMovementTasks = planner.getInterBrokerReplicaMovementTasks(readyBrokers, Collections.emptySet(), _defaultPartitionsMaxCap);
     assertEquals("First task", _partitionMovement0, partitionMovementTasks.get(0).proposal());
-    assertEquals("Second task", _partitionMovement2, partitionMovementTasks.get(1).proposal());
-    assertEquals("Third task", _partitionMovement3, partitionMovementTasks.get(2).proposal());
+    assertEquals("Second task", _partitionMovement3, partitionMovementTasks.get(1).proposal());
+    assertEquals("Third task", _partitionMovement2, partitionMovementTasks.get(2).proposal());
     assertEquals("Fourth task", _partitionMovement1, partitionMovementTasks.get(3).proposal());
   }