text: # 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.
  
  from operator import attrgetter
  
  from kafka.tools import log
  from kafka.tools.assigner.actions import ActionBalanceModule
  
  
  class ActionBalanceCount(ActionBalanceModule):
      name = "count"
      helpstr = "Move the smallest partitions in the cluster to even the partition count per-broker for each replica position"
  
      def process_cluster(self):
          log.info("Starting partition balance by count")
  
          # Figure out the max RF for the cluster and sort all partition lists by size (ascending)
          max_pos = self.cluster.max_replication_factor()
          for broker in self.cluster.brokers:
              for pos in self.cluster.brokers[broker].partitions:
                  self.cluster.brokers[broker].partitions[pos].sort(key=attrgetter('size'))
  
          # Calculate partition counts for each position first
          max_count = {}
          for pos in range(max_pos):
              # Calculate the maximum number of partitions each broker should have (floor(average))
              # We'll also track a remainder and make sure they only go 1 per broker
              pcount = 0
              for broker in self.cluster.brokers:
                  if pos in self.cluster.brokers[broker].partitions:
                      pcount += self.cluster.brokers[broker].num_partitions_at_position(pos)
              max_count[pos] = [pcount / len(self.cluster.brokers), pcount % len(self.cluster.brokers)]
              log.info("Calculating ideal state for replica position {0} - max {1} partitions".format(pos, max_count[pos][0] + 1))
  
          # Balance partition counts for each replica position separately
          for pos in range(max_pos):
              for broker_id in self.cluster.brokers:
                  broker = self.cluster.brokers[broker_id]
                  # Figure out how many more partitions this broker needs
                  diff = max_count[pos][0]
                  if max_count[pos][1]:
                      diff += 1
                      max_count[pos][1] -= 1
                  if pos in broker.partitions:
                      diff -= broker.num_partitions_at_position(pos)
  
                  if diff > 0:
                      log.debug("Moving {0} partitions to broker {1}".format(diff, broker_id))
  
                      # Iterate through the largest brokers to find diff partitions to move to this broker
                      for source_id in self.cluster.brokers:
                          source = self.cluster.brokers[source_id]
                          if diff == 0:
                              break
                          if pos not in source.partitions:
                              continue
  
                          iterlist = list(source.partitions[pos])
                          for partition in iterlist:
                              # If we have moved enough partitions from this broker, exit out of the inner loop
                              if (source.num_partitions_at_position(pos) < max_count[pos][0]) or (diff == 0):
                                  break
  
                              # Skip topics that are being excluded
                              if partition.topic.name in self.args.exclude_topics:
                                  continue
  
                              # If the partition is already on the target, swap positions only if it makes the balance better
                              if broker in partition.replicas:
                                  other_pos = partition.replicas.index(broker)
                                  if (other_pos in source.partitions) and (source.num_partitions_at_position(other_pos) < max_count[other_pos][0]):
                                      partition.swap_replica_positions(source, broker)
                              else:
                                  partition.swap_replicas(source, broker)
                                  diff -= 1
  
                      log.debug("Finish broker {0} with {1} partitions".format(broker_id, broker.num_partitions_at_position(pos)))
                  elif diff < 0:
                      log.debug("Moving {0} partitions off broker {1}".format(-diff, broker_id))
  
                      # Iterate through the smallest brokers to find diff partitions to move off this broker
                      for target_id in self.cluster.brokers:
                          target = self.cluster.brokers[target_id]
                          if diff == 0:
                              break
                          if (pos in target.partitions) and (target.num_partitions_at_position(pos) > (max_count[pos][0] + 1)):
                              continue
  
                          iterlist = list(broker.partitions[pos])
                          for partition in iterlist:
                              # If we have moved enough partitions to this broker, exit out of the inner loop
                              if ((pos in target.partitions) and (target.num_partitions_at_position(pos) >= max_count[pos][0])) or (diff == 0):
                                  break
                              # Skip partitions that are already on the target broker or are being excluded
                              if (target in partition.replicas) or (partition.topic.name in self.args.exclude_topics):
                                  continue
  
                              partition.swap_replicas(broker, target)
                              diff += 1
  
                      log.debug("Finish broker {0} with {1} partitions".format(broker, broker.num_partitions_at_position(pos)))
                  else:
                      log.debug("Skipping broker {0} which has {1} partitions".format(broker, broker.num_partitions_at_position(pos)))
                      continue
  