content: /*
   * 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.clients.consumer.internals;
  
  import org.apache.kafka.clients.consumer.internals.Utils.PartitionComparator;
  import org.apache.kafka.common.Node;
  import org.apache.kafka.common.PartitionInfo;
  import org.apache.kafka.common.TopicPartition;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.Serializable;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.Iterator;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Set;
  import java.util.TreeSet;
  import java.util.stream.Collectors;
  
  /**
   * Sticky assignment implementation used by {@link org.apache.kafka.clients.consumer.StickyAssignor} and
   * {@link org.apache.kafka.clients.consumer.CooperativeStickyAssignor}. Sticky assignors are rack-aware.
   * If racks are specified for consumers, we attempt to match consumer racks with partition replica
   * racks on a best-effort basis, prioritizing balanced assignment over rack alignment. Previously
   * owned partitions may be reassigned to improve rack locality. We use rack-aware assignment if both
   * consumer and partition racks are available and some partitions have replicas only on a subset of racks.
   */
  public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor {
      private static final Logger log = LoggerFactory.getLogger(AbstractStickyAssignor.class);
  
      public static final int DEFAULT_GENERATION = -1;
      private int maxGeneration = DEFAULT_GENERATION;
  
      private PartitionMovements partitionMovements;
  
      // Keep track of the partitions being migrated from one consumer to another during assignment
      // so the cooperative assignor can adjust the assignment
      protected Map<TopicPartition, String> partitionsTransferringOwnership = new HashMap<>();
  
      static final class ConsumerGenerationPair {
          final String consumer;
          final int generation;
          ConsumerGenerationPair(String consumer, int generation) {
              this.consumer = consumer;
              this.generation = generation;
          }
      }
  
      public static final class MemberData {
          public final List<TopicPartition> partitions;
          public final Optional<Integer> generation;
          public final Optional<String> rackId;
  
          public MemberData(List<TopicPartition> partitions, Optional<Integer> generation, Optional<String> rackId) {
              this.partitions = partitions;
              this.generation = generation;
              this.rackId = rackId;
          }
  
          public MemberData(List<TopicPartition> partitions, Optional<Integer> generation) {
              this(partitions, generation, Optional.empty());
          }
      }
  
      protected abstract MemberData memberData(Subscription subscription);
  
      @Override
      public Map<String, List<TopicPartition>> assignPartitions(Map<String, List<PartitionInfo>> partitionsPerTopic,
                                                                Map<String, Subscription> subscriptions) {
          Map<String, List<TopicPartition>> consumerToOwnedPartitions = new HashMap<>();
          Set<TopicPartition> partitionsWithMultiplePreviousOwners = new HashSet<>();
  
          List<PartitionInfo> allPartitions = new ArrayList<>();
          partitionsPerTopic.values().forEach(allPartitions::addAll);
          RackInfo rackInfo = new RackInfo(allPartitions, subscriptions);
  
          AbstractAssignmentBuilder assignmentBuilder;
          if (allSubscriptionsEqual(partitionsPerTopic.keySet(), subscriptions, consumerToOwnedPartitions, partitionsWithMultiplePreviousOwners)) {
              log.debug("Detected that all consumers were subscribed to same set of topics, invoking the "
                            + "optimized assignment algorithm");
              partitionsTransferringOwnership = new HashMap<>();
              assignmentBuilder = new ConstrainedAssignmentBuilder(partitionsPerTopic, rackInfo, consumerToOwnedPartitions, partitionsWithMultiplePreviousOwners);
          } else {
              log.debug("Detected that not all consumers were subscribed to same set of topics, falling back to the "
                            + "general case assignment algorithm");
              // we must set this to null for the general case so the cooperative assignor knows to compute it from scratch
              partitionsTransferringOwnership = null;
              assignmentBuilder = new GeneralAssignmentBuilder(partitionsPerTopic, rackInfo, consumerToOwnedPartitions, subscriptions);
          }
          return assignmentBuilder.build();
      }
  
      public Map<String, List<TopicPartition>> assign(Map<String, Integer> partitionsPerTopic,
                                                      Map<String, Subscription> subscriptions) {
          return assignPartitions(partitionInfosWithoutRacks(partitionsPerTopic), subscriptions);
      }
  
      public int maxGeneration() {
          return maxGeneration;
      }
  
      /**
       * Returns true iff all consumers have an identical subscription. Also fills out the passed in
       * {@code consumerToOwnedPartitions} with each consumer's previously owned and still-subscribed partitions,
       * and the {@code partitionsWithMultiplePreviousOwners} with any partitions claimed by multiple previous owners
       */
      private boolean allSubscriptionsEqual(Set<String> allTopics,
                                            Map<String, Subscription> subscriptions,
                                            Map<String, List<TopicPartition>> consumerToOwnedPartitions,
                                            Set<TopicPartition> partitionsWithMultiplePreviousOwners) {
          boolean isAllSubscriptionsEqual = true;
  
          Set<String> subscribedTopics = new HashSet<>();
  
          // keep track of all previously owned partitions so we can invalidate them if invalid input is
          // detected, eg two consumers somehow claiming the same partition in the same/current generation
          Map<TopicPartition, String> allPreviousPartitionsToOwner = new HashMap<>();
  
          for (Map.Entry<String, Subscription> subscriptionEntry : subscriptions.entrySet()) {
              final String consumer = subscriptionEntry.getKey();
              final Subscription subscription = subscriptionEntry.getValue();
  
              // initialize the subscribed topics set if this is the first subscription
              if (subscribedTopics.isEmpty()) {
                  subscribedTopics.addAll(subscription.topics());
              } else if (isAllSubscriptionsEqual && !(subscription.topics().size() == subscribedTopics.size()
                  && subscribedTopics.containsAll(subscription.topics()))) {
                  isAllSubscriptionsEqual = false;
              }
  
              MemberData memberData = memberData(subscription);
              final int memberGeneration = memberData.generation.orElse(DEFAULT_GENERATION);
              maxGeneration = Math.max(maxGeneration, memberGeneration);
  
              List<TopicPartition> ownedPartitions = new ArrayList<>();
              consumerToOwnedPartitions.put(consumer, ownedPartitions);
  
              // the member has a valid generation, so we can consider its owned partitions if it has the highest
              // generation amongst
              for (final TopicPartition tp : memberData.partitions) {
                  if (allTopics.contains(tp.topic())) {
                      String otherConsumer = allPreviousPartitionsToOwner.get(tp);
                      if (otherConsumer == null) {
                          // this partition is not owned by other consumer in the same generation
                          ownedPartitions.add(tp);
                          allPreviousPartitionsToOwner.put(tp, consumer);
                      } else {
                          final int otherMemberGeneration = subscriptions.get(otherConsumer).generationId().orElse(DEFAULT_GENERATION);
  
                          if (memberGeneration == otherMemberGeneration) {
                              // if two members of the same generation own the same partition, revoke the partition
                              log.error("Found multiple consumers {} and {} claiming the same TopicPartition {} in the "
                                              + "same generation {}, this will be invalidated and removed from their previous assignment.",
                                      consumer, otherConsumer, tp, memberGeneration);
                              partitionsWithMultiplePreviousOwners.add(tp);
                              consumerToOwnedPartitions.get(otherConsumer).remove(tp);
                              allPreviousPartitionsToOwner.put(tp, consumer);
                          } else if (memberGeneration > otherMemberGeneration) {
                              // move partition from the member with an older generation to the member with the newer generation
                              ownedPartitions.add(tp);
                              consumerToOwnedPartitions.get(otherConsumer).remove(tp);
                              allPreviousPartitionsToOwner.put(tp, consumer);
                              log.warn("Consumer {} in generation {} and consumer {} in generation {} claiming the same " +
                                              "TopicPartition {} in different generations. The topic partition will be " +
                                              "assigned to the member with the higher generation {}.",
                                      consumer, memberGeneration,
                                      otherConsumer, otherMemberGeneration,
                                      tp,
                                      memberGeneration);
                          } else {
                              // let the other member continue to own the topic partition
                              log.warn("Consumer {} in generation {} and consumer {} in generation {} claiming the same " +
                                              "TopicPartition {} in different generations. The topic partition will be " +
                                              "assigned to the member with the higher generation {}.",
                                      consumer, memberGeneration,
                                      otherConsumer, otherMemberGeneration,
                                      tp,
                                      otherMemberGeneration);
                          }
                      }
                  }
              }
          }
          return isAllSubscriptionsEqual;
      }
  
      public boolean isSticky() {
          return partitionMovements.isSticky();
      }
  
      public Map<TopicPartition, String> partitionsTransferringOwnership() {
          return partitionsTransferringOwnership;
      }
  
      private static class TopicComparator implements Comparator<String>, Serializable {
          private static final long serialVersionUID = 1L;
          private final Map<String, List<String>> map;
  
          TopicComparator(Map<String, List<String>> map) {
              this.map = map;
          }
  
          @Override
          public int compare(String o1, String o2) {
              int ret = map.get(o1).size() - map.get(o2).size();
              if (ret == 0) {
                  ret = o1.compareTo(o2);
              }
              return ret;
          }
      }
  
      private static class SubscriptionComparator implements Comparator<String>, Serializable {
          private static final long serialVersionUID = 1L;
          private final Map<String, List<TopicPartition>> map;
  
          SubscriptionComparator(Map<String, List<TopicPartition>> map) {
              this.map = map;
          }
  
          @Override
          public int compare(String o1, String o2) {
              int ret = map.get(o1).size() - map.get(o2).size();
              if (ret == 0)
                  ret = o1.compareTo(o2);
              return ret;
          }
      }
  
      /**
       * This class maintains some data structures to simplify lookup of partition movements among consumers. At each point of
       * time during a partition rebalance it keeps track of partition movements corresponding to each topic, and also possible
       * movement (in form a <code>ConsumerPair</code> object) for each partition.
       */
      private static class PartitionMovements {
          private final Map<String, Map<ConsumerPair, Set<TopicPartition>>> partitionMovementsByTopic = new HashMap<>();
          private final Map<TopicPartition, ConsumerPair> partitionMovements = new HashMap<>();
  
          private ConsumerPair removeMovementRecordOfPartition(TopicPartition partition) {
              ConsumerPair pair = partitionMovements.remove(partition);
  
              String topic = partition.topic();
              Map<ConsumerPair, Set<TopicPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
              partitionMovementsForThisTopic.get(pair).remove(partition);
              if (partitionMovementsForThisTopic.get(pair).isEmpty())
                  partitionMovementsForThisTopic.remove(pair);
              if (partitionMovementsByTopic.get(topic).isEmpty())
                  partitionMovementsByTopic.remove(topic);
  
              return pair;
          }
  
          private void addPartitionMovementRecord(TopicPartition partition, ConsumerPair pair) {
              partitionMovements.put(partition, pair);
  
              String topic = partition.topic();
              if (!partitionMovementsByTopic.containsKey(topic))
                  partitionMovementsByTopic.put(topic, new HashMap<>());
  
              Map<ConsumerPair, Set<TopicPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
              if (!partitionMovementsForThisTopic.containsKey(pair))
                  partitionMovementsForThisTopic.put(pair, new HashSet<>());
  
              partitionMovementsForThisTopic.get(pair).add(partition);
          }
  
          private void movePartition(TopicPartition partition, String oldConsumer, String newConsumer) {
              ConsumerPair pair = new ConsumerPair(oldConsumer, newConsumer);
  
              if (partitionMovements.containsKey(partition)) {
                  // this partition has previously moved
                  ConsumerPair existingPair = removeMovementRecordOfPartition(partition);
                  assert existingPair.dstMemberId.equals(oldConsumer);
                  if (!existingPair.srcMemberId.equals(newConsumer)) {
                      // the partition is not moving back to its previous consumer
                      // return new ConsumerPair2(existingPair.src, newConsumer);
                      addPartitionMovementRecord(partition, new ConsumerPair(existingPair.srcMemberId, newConsumer));
                  }
              } else
                  addPartitionMovementRecord(partition, pair);
          }
  
          private TopicPartition getTheActualPartitionToBeMoved(TopicPartition partition, String oldConsumer, String newConsumer) {
              String topic = partition.topic();
  
              if (!partitionMovementsByTopic.containsKey(topic))
                  return partition;
  
              if (partitionMovements.containsKey(partition)) {
                  // this partition has previously moved
                  assert oldConsumer.equals(partitionMovements.get(partition).dstMemberId);
                  oldConsumer = partitionMovements.get(partition).srcMemberId;
              }
  
              Map<ConsumerPair, Set<TopicPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
              ConsumerPair reversePair = new ConsumerPair(newConsumer, oldConsumer);
              if (!partitionMovementsForThisTopic.containsKey(reversePair))
                  return partition;
  
              return partitionMovementsForThisTopic.get(reversePair).iterator().next();
          }
  
          private boolean isLinked(String src, String dst, Set<ConsumerPair> pairs, List<String> currentPath) {
              if (src.equals(dst))
                  return false;
  
              if (pairs.isEmpty())
                  return false;
  
              if (new ConsumerPair(src, dst).in(pairs)) {
                  currentPath.add(src);
                  currentPath.add(dst);
                  return true;
              }
  
              for (ConsumerPair pair: pairs)
                  if (pair.srcMemberId.equals(src)) {
                      Set<ConsumerPair> reducedSet = new HashSet<>(pairs);
                      reducedSet.remove(pair);
                      currentPath.add(pair.srcMemberId);
                      return isLinked(pair.dstMemberId, dst, reducedSet, currentPath);
                  }
  
              return false;
          }
  
          private boolean in(List<String> cycle, Set<List<String>> cycles) {
              List<String> superCycle = new ArrayList<>(cycle);
              superCycle.remove(superCycle.size() - 1);
              superCycle.addAll(cycle);
              for (List<String> foundCycle: cycles) {
                  if (foundCycle.size() == cycle.size() && Collections.indexOfSubList(superCycle, foundCycle) != -1)
                      return true;
              }
              return false;
          }
  
          private boolean hasCycles(Set<ConsumerPair> pairs) {
              Set<List<String>> cycles = new HashSet<>();
              for (ConsumerPair pair: pairs) {
                  Set<ConsumerPair> reducedPairs = new HashSet<>(pairs);
                  reducedPairs.remove(pair);
                  List<String> path = new ArrayList<>(Collections.singleton(pair.srcMemberId));
                  if (isLinked(pair.dstMemberId, pair.srcMemberId, reducedPairs, path) && !in(path, cycles)) {
                      cycles.add(new ArrayList<>(path));
                      log.error("A cycle of length {} was found: {}", path.size() - 1, path);
                  }
              }
  
              // for now we want to make sure there is no partition movements of the same topic between a pair of consumers.
              // the odds of finding a cycle among more than two consumers seem to be very low (according to various randomized
              // tests with the given sticky algorithm) that it should not worth the added complexity of handling those cases.
              for (List<String> cycle: cycles)
                  if (cycle.size() == 3) // indicates a cycle of length 2
                      return true;
              return false;
          }
  
          private boolean isSticky() {
              for (Map.Entry<String, Map<ConsumerPair, Set<TopicPartition>>> topicMovements: this.partitionMovementsByTopic.entrySet()) {
                  Set<ConsumerPair> topicMovementPairs = topicMovements.getValue().keySet();
                  if (hasCycles(topicMovementPairs)) {
                      log.error("Stickiness is violated for topic {}"
                          + "\nPartition movements for this topic occurred among the following consumer pairs:"
                          + "\n{}", topicMovements.getKey(), topicMovements.getValue().toString());
                      return false;
                  }
              }
  
              return true;
          }
      }
  
      /**
       * <code>ConsumerPair</code> represents a pair of Kafka consumer ids involved in a partition reassignment. Each
       * <code>ConsumerPair</code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
       * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
       * partition of the particular topic was moved from the source consumer to the destination consumer during the rebalance.
       * This class is used, through the <code>PartitionMovements</code> class, by the sticky assignor and helps in determining
       * whether a partition reassignment results in cycles among the generated graph of consumer pairs.
       */
      private static class ConsumerPair {
          private final String srcMemberId;
          private final String dstMemberId;
  
          ConsumerPair(String srcMemberId, String dstMemberId) {
              this.srcMemberId = srcMemberId;
              this.dstMemberId = dstMemberId;
          }
  
          public String toString() {
              return this.srcMemberId + "->" + this.dstMemberId;
          }
  
          @Override
          public int hashCode() {
              final int prime = 31;
              int result = 1;
              result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
              result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
              return result;
          }
  
          @Override
          public boolean equals(Object obj) {
              if (obj == null)
                  return false;
  
              if (!getClass().isInstance(obj))
                  return false;
  
              ConsumerPair otherPair = (ConsumerPair) obj;
              return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
          }
  
          private boolean in(Set<ConsumerPair> pairs) {
              for (ConsumerPair pair: pairs)
                  if (this.equals(pair))
                      return true;
              return false;
          }
      }
  
      private class RackInfo {
          private final Map<String, String> consumerRacks;
          private final Map<TopicPartition, Set<String>> partitionRacks;
          private final Map<TopicPartition, Integer> numConsumersByPartition;
  
          public RackInfo(List<PartitionInfo> partitionInfos, Map<String, Subscription> subscriptions) {
              List<Subscription> consumers = new ArrayList<>(subscriptions.values());
  
              Map<String, List<String>> consumersByRack = new HashMap<>();
              subscriptions.forEach((memberId, subscription) ->
                      subscription.rackId().filter(r -> !r.isEmpty()).ifPresent(rackId -> put(consumersByRack, rackId, memberId)));
  
              Map<String, List<TopicPartition>> partitionsByRack;
              Map<TopicPartition, Set<String>> partitionRacks;
              if (consumersByRack.isEmpty()) {
                  partitionsByRack = Collections.emptyMap();
                  partitionRacks = Collections.emptyMap();
              } else {
                  partitionRacks = new HashMap<>(partitionInfos.size());
                  partitionsByRack = new HashMap<>();
                  partitionInfos.forEach(p -> {
                      TopicPartition tp = new TopicPartition(p.topic(), p.partition());
                      Set<String> racks = new HashSet<>(p.replicas().length);
                      partitionRacks.put(tp, racks);
                      Arrays.stream(p.replicas())
                              .map(Node::rack)
                              .filter(Objects::nonNull)
                              .distinct()
                              .forEach(rackId -> {
                                  put(partitionsByRack, rackId, tp);
                                  racks.add(rackId);
                              });
                  });
              }
  
              if (useRackAwareAssignment(consumersByRack.keySet(), partitionsByRack.keySet(), partitionRacks)) {
                  this.consumerRacks = new HashMap<>(consumers.size());
                  consumersByRack.forEach((rack, rackConsumers) -> rackConsumers.forEach(c -> consumerRacks.put(c, rack)));
                  this.partitionRacks = partitionRacks;
              } else {
                  this.consumerRacks = Collections.emptyMap();
                  this.partitionRacks = Collections.emptyMap();
              }
              numConsumersByPartition = partitionRacks.entrySet().stream()
                      .collect(Collectors.toMap(Entry::getKey, e -> e.getValue().stream()
                          .map(r -> consumersByRack.getOrDefault(r, Collections.emptyList()).size())
                          .reduce(0, Integer::sum)));
          }
  
          private boolean racksMismatch(String consumer, TopicPartition tp) {
              String consumerRack = consumerRacks.get(consumer);
              Set<String> replicaRacks = partitionRacks.get(tp);
              return consumerRack != null && (replicaRacks == null || !replicaRacks.contains(consumerRack));
          }
  
          private List<TopicPartition> sortPartitionsByRackConsumers(List<TopicPartition> partitions) {
              if (numConsumersByPartition.isEmpty())
                  return partitions;
              // Return a sorted linked list of partitions to enable fast updates during rack-aware assignment
              List<TopicPartition> sortedPartitions = new LinkedList<>(partitions);
              sortedPartitions.sort(Comparator.comparing(tp -> numConsumersByPartition.getOrDefault(tp, 0)));
              return sortedPartitions;
          }
  
          private int nextRackConsumer(TopicPartition tp, List<String> consumerList, int firstIndex) {
              Set<String> racks = partitionRacks.get(tp);
              if (racks == null || racks.isEmpty())
                  return -1;
              for (int i = 0; i < consumerList.size(); i++) {
                  int index = (firstIndex + i) % consumerList.size();
                  String consumer = consumerList.get(index);
                  String consumerRack = consumerRacks.get(consumer);
                  if (consumerRack != null && racks.contains(consumerRack))
                      return index;
              }
              return -1;
          }
  
          @Override
          public String toString() {
              return "RackInfo(" +
                      "consumerRacks=" + consumerRacks +
                      ", partitionRacks=" + partitionRacks +
                      ")";
          }
      }
  
      private abstract class AbstractAssignmentBuilder {
  
          final Map<String, List<PartitionInfo>> partitionsPerTopic;
          final RackInfo rackInfo;
          final Map<String, List<TopicPartition>> currentAssignment;
          final int totalPartitionsCount;
  
          AbstractAssignmentBuilder(Map<String, List<PartitionInfo>> partitionsPerTopic,
                                    RackInfo rackInfo,
                                    Map<String, List<TopicPartition>> currentAssignment) {
              this.partitionsPerTopic = partitionsPerTopic;
              this.currentAssignment = currentAssignment;
              this.rackInfo = rackInfo;
              this.totalPartitionsCount = partitionsPerTopic.values().stream().map(List::size).reduce(0, Integer::sum);
          }
  
          /**
           * Builds the assignment.
           *
           * @return Map from each member to the list of partitions assigned to them.
           */
          abstract Map<String, List<TopicPartition>> build();
  
          protected List<TopicPartition> getAllTopicPartitions(List<String> sortedAllTopics) {
              List<TopicPartition> allPartitions = new ArrayList<>(totalPartitionsCount);
  
              for (String topic : sortedAllTopics) {
                  partitionsPerTopic.get(topic).forEach(p -> allPartitions.add(new TopicPartition(p.topic(), p.partition())));
              }
              return allPartitions;
          }
      }
  
      /**
       * This constrained assignment optimizes the assignment algorithm when all consumers were subscribed to same set of topics.
       * The method includes the following steps:
       *
       * 1. Reassign previously owned partitions:
       *   a. if owned less than minQuota partitions, just assign all owned partitions, and put the member into unfilled member list.
       *   b. if owned maxQuota or more, and we're still under the number of expected max capacity members, assign maxQuota partitions
       *   c. if owned at least "minQuota" of partitions, assign minQuota partitions, and put the member into unfilled member list if
       *     we're still under the number of expected max capacity members
       *   If using rack-aware algorithm, only owned partitions with matching racks are allocated in this step.
       * 2. Fill remaining members with rack matching up to the expected numbers of maxQuota partitions, otherwise, to minQuota partitions.
       *    Partitions that cannot be aligned on racks within the quota are not assigned in this step. This step is only used if rack-aware.
       * 3. Fill remaining members up to the expected numbers of maxQuota partitions, otherwise, to minQuota partitions.
       *    For rack-aware algorithm, these are partitions that could not be aligned on racks within the balancing constraints.
       *
       */
      private class ConstrainedAssignmentBuilder extends AbstractAssignmentBuilder {
  
          private final Set<TopicPartition> partitionsWithMultiplePreviousOwners;
          private final Map<TopicPartition, String> maybeRevokedPartitions;
  
          // the consumers which may still be assigned one or more partitions to reach expected capacity
          private final List<String> unfilledMembersWithUnderMinQuotaPartitions;
          private final LinkedList<String> unfilledMembersWithExactlyMinQuotaPartitions;
  
          private final int minQuota;
          private final int maxQuota;
          // the expected number of members receiving more than minQuota partitions (zero when minQuota == maxQuota)
          private final int expectedNumMembersWithOverMinQuotaPartitions;
          // the current number of members receiving more than minQuota partitions (zero when minQuota == maxQuota)
          private int currentNumMembersWithOverMinQuotaPartitions;
  
          private final Map<String, List<TopicPartition>> assignment;
          private final List<TopicPartition> assignedPartitions;
  
          /**
           * Constructs a constrained assignment builder.
           *
           * @param partitionsPerTopic                   The partitions for each subscribed topic
           * @param rackInfo                             Rack information for consumers and racks
           * @param consumerToOwnedPartitions            Each consumer's previously owned and still-subscribed partitions
           * @param partitionsWithMultiplePreviousOwners The partitions being claimed in the previous assignment of multiple consumers
           */
          ConstrainedAssignmentBuilder(Map<String, List<PartitionInfo>> partitionsPerTopic,
                                       RackInfo rackInfo,
                                       Map<String, List<TopicPartition>> consumerToOwnedPartitions,
                                       Set<TopicPartition> partitionsWithMultiplePreviousOwners) {
              super(partitionsPerTopic, rackInfo, consumerToOwnedPartitions);
  
              this.partitionsWithMultiplePreviousOwners = partitionsWithMultiplePreviousOwners;
              maybeRevokedPartitions = new HashMap<>();
              unfilledMembersWithUnderMinQuotaPartitions = new LinkedList<>();
              unfilledMembersWithExactlyMinQuotaPartitions = new LinkedList<>();
  
              int numberOfConsumers = consumerToOwnedPartitions.size();
  
              minQuota = (int) Math.floor(((double) totalPartitionsCount) / numberOfConsumers);
              maxQuota = (int) Math.ceil(((double) totalPartitionsCount) / numberOfConsumers);
              expectedNumMembersWithOverMinQuotaPartitions = totalPartitionsCount % numberOfConsumers;
              currentNumMembersWithOverMinQuotaPartitions = 0;
  
              // initialize the assignment map with an empty array of size maxQuota for all members
              assignment = new HashMap<>(consumerToOwnedPartitions.keySet().stream()
                      .collect(Collectors.toMap(c -> c, c -> new ArrayList<>(maxQuota))));
              assignedPartitions = new ArrayList<>();
          }
  
          @Override
          Map<String, List<TopicPartition>> build() {
              if (log.isDebugEnabled()) {
                  log.debug("Performing constrained assign with partitionsPerTopic: {}, currentAssignment: {}, rackInfo {}.",
                          partitionsPerTopic, currentAssignment, rackInfo);
              }
  
              assignOwnedPartitions();
  
              List<TopicPartition> unassignedPartitions = getUnassignedPartitions(assignedPartitions);
  
              if (log.isDebugEnabled()) {
                  log.debug("After reassigning previously owned partitions, unfilled members: {}, unassigned partitions: {}, " +
                          "current assignment: {}", unfilledMembersWithUnderMinQuotaPartitions, unassignedPartitions, assignment);
              }
  
              Collections.sort(unfilledMembersWithUnderMinQuotaPartitions);
              Collections.sort(unfilledMembersWithExactlyMinQuotaPartitions);
              unassignedPartitions = rackInfo.sortPartitionsByRackConsumers(unassignedPartitions);
  
              assignRackAwareRoundRobin(unassignedPartitions);
              assignRoundRobin(unassignedPartitions);
              verifyUnfilledMembers();
  
              log.info("Final assignment of partitions to consumers: \n{}", assignment);
  
              return assignment;
          }
  
          // Reassign previously owned partitions, up to the expected number of partitions per consumer
          private void assignOwnedPartitions() {
  
              for (Map.Entry<String, List<TopicPartition>> consumerEntry : currentAssignment.entrySet()) {
                  String consumer = consumerEntry.getKey();
                  List<TopicPartition> ownedPartitions = consumerEntry.getValue().stream()
                          .filter(tp -> {
                              boolean mismatch = rackInfo.racksMismatch(consumer, tp);
                              if (mismatch) {
                                  maybeRevokedPartitions.put(tp, consumer);
                              }
                              return !mismatch;
                          })
                          .sorted(Comparator.comparing(TopicPartition::partition).thenComparing(TopicPartition::topic))
                          .collect(Collectors.toList());
  
                  List<TopicPartition> consumerAssignment = assignment.get(consumer);
  
                  for (TopicPartition doublyClaimedPartition : partitionsWithMultiplePreviousOwners) {
                      if (ownedPartitions.contains(doublyClaimedPartition)) {
                          log.error("Found partition {} still claimed as owned by consumer {}, despite being claimed by multiple "
                                          + "consumers already in the same generation. Removing it from the ownedPartitions",
                                  doublyClaimedPartition, consumer);
                          ownedPartitions.remove(doublyClaimedPartition);
                      }
                  }
  
                  if (ownedPartitions.size() < minQuota) {
                      // the expected assignment size is more than this consumer has now, so keep all the owned partitions
                      // and put this member into the unfilled member list
                      if (ownedPartitions.size() > 0) {
                          consumerAssignment.addAll(ownedPartitions);
                          assignedPartitions.addAll(ownedPartitions);
                      }
                      unfilledMembersWithUnderMinQuotaPartitions.add(consumer);
                  } else if (ownedPartitions.size() >= maxQuota && currentNumMembersWithOverMinQuotaPartitions < expectedNumMembersWithOverMinQuotaPartitions) {
                      // consumer owned the "maxQuota" of partitions or more, and we're still under the number of expected members
                      // with more than the minQuota partitions, so keep "maxQuota" of the owned partitions, and revoke the rest of the partitions
                      currentNumMembersWithOverMinQuotaPartitions++;
                      if (currentNumMembersWithOverMinQuotaPartitions == expectedNumMembersWithOverMinQuotaPartitions) {
                          unfilledMembersWithExactlyMinQuotaPartitions.clear();
                      }
                      List<TopicPartition> maxQuotaPartitions = ownedPartitions.subList(0, maxQuota);
                      consumerAssignment.addAll(maxQuotaPartitions);
                      assignedPartitions.addAll(maxQuotaPartitions);
                      for (TopicPartition topicPartition : ownedPartitions.subList(maxQuota, ownedPartitions.size())) {
                          maybeRevokedPartitions.put(topicPartition, consumer);
                      }
                  } else {
                      // consumer owned at least "minQuota" of partitions
                      // so keep "minQuota" of the owned partitions, and revoke the rest of the partitions
                      List<TopicPartition> minQuotaPartitions = ownedPartitions.subList(0, minQuota);
                      consumerAssignment.addAll(minQuotaPartitions);
                      assignedPartitions.addAll(minQuotaPartitions);
                      for (TopicPartition topicPartition : ownedPartitions.subList(minQuota, ownedPartitions.size())) {
                          maybeRevokedPartitions.put(topicPartition, consumer);
                      }
                      // this consumer is potential maxQuota candidate since we're still under the number of expected members
                      // with more than the minQuota partitions. Note, if the number of expected members with more than
                      // the minQuota partitions is 0, it means minQuota == maxQuota, and there are no potentially unfilled
                      if (currentNumMembersWithOverMinQuotaPartitions < expectedNumMembersWithOverMinQuotaPartitions) {
                          unfilledMembersWithExactlyMinQuotaPartitions.add(consumer);
                      }
                  }
              }
          }
  
          // Round-Robin filling within racks for remaining members up to the expected numbers of maxQuota,
          // otherwise, to minQuota
          private void assignRackAwareRoundRobin(List<TopicPartition> unassignedPartitions) {
              if (rackInfo.consumerRacks.isEmpty())
                  return;
              int nextUnfilledConsumerIndex = 0;
              Iterator<TopicPartition> unassignedIter = unassignedPartitions.iterator();
              while (unassignedIter.hasNext()) {
                  TopicPartition unassignedPartition = unassignedIter.next();
                  String consumer = null;
                  int nextIndex = rackInfo.nextRackConsumer(unassignedPartition, unfilledMembersWithUnderMinQuotaPartitions, nextUnfilledConsumerIndex);
                  if (nextIndex >= 0) {
                      consumer = unfilledMembersWithUnderMinQuotaPartitions.get(nextIndex);
                      int assignmentCount = assignment.get(consumer).size() + 1;
                      if (assignmentCount >= minQuota) {
                          unfilledMembersWithUnderMinQuotaPartitions.remove(consumer);
                          // Only add this consumer if the current num members at maxQuota is less than the expected number
                          // since a consumer at minQuota can only be considered unfilled if it's possible to add another partition,
                          // which would bump it to maxQuota and exceed the expectedNumMembersWithOverMinQuotaPartitions
                          if (assignmentCount < maxQuota && (currentNumMembersWithOverMinQuotaPartitions < expectedNumMembersWithOverMinQuotaPartitions)) {
                              unfilledMembersWithExactlyMinQuotaPartitions.add(consumer);
                          }
                      } else {
                          nextIndex++;
                      }
                      nextUnfilledConsumerIndex = unfilledMembersWithUnderMinQuotaPartitions.isEmpty() ? 0 : nextIndex % unfilledMembersWithUnderMinQuotaPartitions.size();
                  } else if (!unfilledMembersWithExactlyMinQuotaPartitions.isEmpty()) {
                      int firstIndex = rackInfo.nextRackConsumer(unassignedPartition, unfilledMembersWithExactlyMinQuotaPartitions, 0);
                      if (firstIndex >= 0) {
                          consumer = unfilledMembersWithExactlyMinQuotaPartitions.get(firstIndex);
                          if (assignment.get(consumer).size() + 1 == maxQuota) {
                              unfilledMembersWithExactlyMinQuotaPartitions.remove(firstIndex);
                              currentNumMembersWithOverMinQuotaPartitions++;
                              // Clear this once the current num consumers over minQuota reaches the expected number since this
                              // means all consumers at minQuota are now considered filled
                              if (currentNumMembersWithOverMinQuotaPartitions == expectedNumMembersWithOverMinQuotaPartitions) {
                                  unfilledMembersWithExactlyMinQuotaPartitions.clear();
                              }
                          }
                      }
                  }
  
                  if (consumer != null) {
                      assignNewPartition(unassignedPartition, consumer);
                      unassignedIter.remove();
                  }
              }
          }
  
          private void assignRoundRobin(List<TopicPartition> unassignedPartitions) {
  
              Iterator<String> unfilledConsumerIter = unfilledMembersWithUnderMinQuotaPartitions.iterator();
              // Round-Robin filling remaining members up to the expected numbers of maxQuota, otherwise, to minQuota
              for (TopicPartition unassignedPartition : unassignedPartitions) {
                  String consumer;
                  if (unfilledConsumerIter.hasNext()) {
                      consumer = unfilledConsumerIter.next();
                  } else {
                      if (unfilledMembersWithUnderMinQuotaPartitions.isEmpty() && unfilledMembersWithExactlyMinQuotaPartitions.isEmpty()) {
                          // Should not enter here since we have calculated the exact number to assign to each consumer.
                          // This indicates issues in the assignment algorithm
                          int currentPartitionIndex = unassignedPartitions.indexOf(unassignedPartition);
                          log.error("No more unfilled consumers to be assigned. The remaining unassigned partitions are: {}",
                                  unassignedPartitions.subList(currentPartitionIndex, unassignedPartitions.size()));
                          throw new IllegalStateException("No more unfilled consumers to be assigned.");
                      } else if (unfilledMembersWithUnderMinQuotaPartitions.isEmpty()) {
                          consumer = unfilledMembersWithExactlyMinQuotaPartitions.poll();
                      } else {
                          unfilledConsumerIter = unfilledMembersWithUnderMinQuotaPartitions.iterator();
                          consumer = unfilledConsumerIter.next();
                      }
                  }
  
                  int currentAssignedCount = assignNewPartition(unassignedPartition, consumer);
  
                  if (currentAssignedCount == minQuota) {
                      unfilledConsumerIter.remove();
                      unfilledMembersWithExactlyMinQuotaPartitions.add(consumer);
                  } else if (currentAssignedCount == maxQuota) {
                      currentNumMembersWithOverMinQuotaPartitions++;
                      if (currentNumMembersWithOverMinQuotaPartitions == expectedNumMembersWithOverMinQuotaPartitions) {
                          // We only start to iterate over the "potentially unfilled" members at minQuota after we've filled
                          // all members up to at least minQuota, so once the last minQuota member reaches maxQuota, we
                          // should be done. But in case of some algorithmic error, just log a warning and continue to
                          // assign any remaining partitions within the assignment constraints
                          if (unassignedPartitions.indexOf(unassignedPartition) != unassignedPartitions.size() - 1) {
                              log.error("Filled the last member up to maxQuota but still had partitions remaining to assign, "
                                      + "will continue but this indicates a bug in the assignment.");
                          }
                      }
                  }
              }
          }
  
          private int assignNewPartition(TopicPartition unassignedPartition, String consumer) {
              List<TopicPartition> consumerAssignment = assignment.get(consumer);
              consumerAssignment.add(unassignedPartition);
  
              // We already assigned all possible ownedPartitions, so we know this must be newly assigned to this consumer
              // or else the partition was actually claimed by multiple previous owners and had to be invalidated from all
              // members claimed ownedPartitions
              if ((maybeRevokedPartitions.containsKey(unassignedPartition) && !maybeRevokedPartitions.get(unassignedPartition).equals(consumer))
                      || partitionsWithMultiplePreviousOwners.contains(unassignedPartition)) {
                  partitionsTransferringOwnership.put(unassignedPartition, consumer);
              }
  
              return consumerAssignment.size();
          }
  
          private void verifyUnfilledMembers() {
  
              if (!unfilledMembersWithUnderMinQuotaPartitions.isEmpty()) {
                  // we expected all the remaining unfilled members have minQuota partitions and we're already at the expected number
                  // of members with more than the minQuota partitions. Otherwise, there must be error here.
                  if (currentNumMembersWithOverMinQuotaPartitions != expectedNumMembersWithOverMinQuotaPartitions) {
                      log.error("Current number of members with more than the minQuota partitions: {}, is less than the expected number " +
                                      "of members with more than the minQuota partitions: {}, and no more partitions to be assigned to the remaining unfilled consumers: {}",
                              currentNumMembersWithOverMinQuotaPartitions, expectedNumMembersWithOverMinQuotaPartitions, unfilledMembersWithUnderMinQuotaPartitions);
                      throw new IllegalStateException("We haven't reached the expected number of members with " +
                              "more than the minQuota partitions, but no more partitions to be assigned");
                  } else {
                      for (String unfilledMember : unfilledMembersWithUnderMinQuotaPartitions) {
                          int assignedPartitionsCount = assignment.get(unfilledMember).size();
                          if (assignedPartitionsCount != minQuota) {
                              log.error("Consumer: [{}] should have {} partitions, but got {} partitions, and no more partitions " +
                                      "to be assigned. The remaining unfilled consumers are: {}", unfilledMember, minQuota, assignedPartitionsCount, unfilledMembersWithUnderMinQuotaPartitions);
                              throw new IllegalStateException(String.format("Consumer: [%s] doesn't reach minQuota partitions, " +
                                      "and no more partitions to be assigned", unfilledMember));
                          } else {
                              log.trace("skip over this unfilled member: [{}] because we've reached the expected number of " +
                                      "members with more than the minQuota partitions, and this member already has minQuota partitions", unfilledMember);
                          }
                      }
                  }
              }
          }
  
          /**
           * get the unassigned partition list by computing the difference set of all sorted partitions
           * and sortedAssignedPartitions. If no assigned partitions, we'll just return all sorted topic partitions.
           *
           * To compute the difference set, we use two pointers technique here:
           *
           * We loop through the all sorted topics, and then iterate all partitions the topic has,
           * compared with the ith element in sortedAssignedPartitions(i starts from 0):
           *   - if not equal to the ith element, add to unassignedPartitions
           *   - if equal to the ith element, get next element from sortedAssignedPartitions
           *
           * @param sortedAssignedPartitions  sorted partitions, all are included in the sortedPartitions
           * @return                          the partitions not yet assigned to any consumers
           */
          private List<TopicPartition> getUnassignedPartitions(List<TopicPartition> sortedAssignedPartitions) {
              List<String> sortedAllTopics = new ArrayList<>(partitionsPerTopic.keySet());
              // sort all topics first, then we can have sorted all topic partitions by adding partitions starting from 0
              Collections.sort(sortedAllTopics);
  
              if (sortedAssignedPartitions.isEmpty()) {
                  // no assigned partitions means all partitions are unassigned partitions
                  return getAllTopicPartitions(sortedAllTopics);
              }
  
              List<TopicPartition> unassignedPartitions = new ArrayList<>(totalPartitionsCount - sortedAssignedPartitions.size());
  
              sortedAssignedPartitions.sort(Comparator.comparing(TopicPartition::topic).thenComparing(TopicPartition::partition));
  
              boolean shouldAddDirectly = false;
              Iterator<TopicPartition> sortedAssignedPartitionsIter = sortedAssignedPartitions.iterator();
              TopicPartition nextAssignedPartition = sortedAssignedPartitionsIter.next();
  
              for (String topic : sortedAllTopics) {
                  int partitionCount = partitionsPerTopic.get(topic).size();
                  for (int i = 0; i < partitionCount; i++) {
                      if (shouldAddDirectly || !(nextAssignedPartition.topic().equals(topic) && nextAssignedPartition.partition() == i)) {
                          unassignedPartitions.add(new TopicPartition(topic, i));
                      } else {
                          // this partition is in assignedPartitions, don't add to unassignedPartitions, just get next assigned partition
                          if (sortedAssignedPartitionsIter.hasNext()) {
                              nextAssignedPartition = sortedAssignedPartitionsIter.next();
                          } else {
                              // add the remaining directly since there is no more sortedAssignedPartitions
                              shouldAddDirectly = true;
                          }
                      }
                  }
              }
  
              return unassignedPartitions;
          }
      }
  
      /**
       * This general assignment algorithm guarantees the assignment that is as balanced as possible.
       * This method includes the following steps:
       *
       * 1. Preserving all the existing partition assignments. If rack-aware algorithm is used, only assignments
       *    within racks are preserved.
       * 2. Removing all the partition assignments that have become invalid due to the change that triggers the reassignment.
       *    Partition assignments with mismatched racks are also removed.
       * 3. Assigning the unassigned partitions in a way that balances out the overall assignments of partitions to consumers.
       *    while preserving rack-alignment. This step is used only for rack-aware assignment.
       * 4. Assigning the remaining unassigned partitions in a way that balances out the overall assignments of partitions to consumers.
       *    For rack-aware algorithm, these are partitions that could not be aligned on racks within the balancing constraints.
       * 5. Further balancing out the resulting assignment by finding the partitions that can be reassigned
       *    to another consumer towards an overall more balanced assignment. For rack-aware algorithm, attempt
       *    to retain rack alignment if possible.
       *
       */
      private class GeneralAssignmentBuilder extends AbstractAssignmentBuilder {
          private final Map<String, Subscription> subscriptions;
  
          // a mapping of all topics to all consumers that can be assigned to them
          private final Map<String, List<String>> topic2AllPotentialConsumers;
          // a mapping of all consumers to all potential topics that can be assigned to them
          private final Map<String, List<String>> consumer2AllPotentialTopics;
          // a mapping of partition to current consumer
          private final Map<TopicPartition, String> currentPartitionConsumer;
          private final List<TopicPartition> sortedAllPartitions;
          // an ascending sorted set of consumers based on how many topic partitions are already assigned to them
          private final TreeSet<String> sortedCurrentSubscriptions;
          private boolean revocationRequired;
  
          /**
           * Constructs a general assignment builder.
           *
           * @param partitionsPerTopic         The partitions for each subscribed topic.
           * @param subscriptions              Map from the member id to their respective topic subscription
           * @param currentAssignment          Each consumer's previously owned and still-subscribed partitions
           * @param rackInfo                   Rack information for consumers and partitions
           */
          GeneralAssignmentBuilder(Map<String, List<PartitionInfo>> partitionsPerTopic,
                                   RackInfo rackInfo,
                                   Map<String, List<TopicPartition>> currentAssignment,
                                   Map<String, Subscription> subscriptions) {
              super(partitionsPerTopic, rackInfo, currentAssignment);
              this.subscriptions = subscriptions;
  
              topic2AllPotentialConsumers = new HashMap<>(partitionsPerTopic.keySet().size());
              consumer2AllPotentialTopics = new HashMap<>(subscriptions.keySet().size());
  
              // initialize topic2AllPotentialConsumers and consumer2AllPotentialTopics
              partitionsPerTopic.keySet().forEach(
                      topicName -> topic2AllPotentialConsumers.put(topicName, new ArrayList<>()));
  
              subscriptions.forEach((consumerId, subscription) -> {
                  List<String> subscribedTopics = new ArrayList<>(subscription.topics().size());
                  consumer2AllPotentialTopics.put(consumerId, subscribedTopics);
                  subscription.topics().stream().filter(topic -> partitionsPerTopic.get(topic) != null).forEach(topic -> {
                      subscribedTopics.add(topic);
                      topic2AllPotentialConsumers.get(topic).add(consumerId);
                  });
  
                  // add this consumer to currentAssignment (with an empty topic partition assignment) if it does not already exist
                  if (!currentAssignment.containsKey(consumerId))
                      currentAssignment.put(consumerId, new ArrayList<>());
              });
  
              currentPartitionConsumer = new HashMap<>();
              for (Map.Entry<String, List<TopicPartition>> entry: currentAssignment.entrySet())
                  for (TopicPartition topicPartition: entry.getValue())
                      currentPartitionConsumer.put(topicPartition, entry.getKey());
  
              List<String> sortedAllTopics = new ArrayList<>(topic2AllPotentialConsumers.keySet());
              sortedAllTopics.sort(new TopicComparator(topic2AllPotentialConsumers));
              sortedAllPartitions = getAllTopicPartitions(sortedAllTopics);
  
              sortedCurrentSubscriptions = new TreeSet<>(new SubscriptionComparator(currentAssignment));
          }
  
          @Override
          Map<String, List<TopicPartition>> build() {
              if (log.isDebugEnabled()) {
                  log.debug("performing general assign. partitionsPerTopic: {}, subscriptions: {}, currentAssignment: {}, rackInfo: {}",
                          partitionsPerTopic, subscriptions, currentAssignment, rackInfo);
              }
  
              Map<TopicPartition, ConsumerGenerationPair> prevAssignment = new HashMap<>();
              partitionMovements = new PartitionMovements();
              prepopulateCurrentAssignments(prevAssignment);
  
              // the partitions already assigned in current assignment
              List<TopicPartition> assignedPartitions = assignOwnedPartitions();
  
              // all partitions that still need to be assigned
              List<TopicPartition> unassignedPartitions = getUnassignedPartitions(assignedPartitions);
  
              if (log.isDebugEnabled()) {
                  log.debug("unassigned Partitions: {}", unassignedPartitions);
              }
  
              // at this point we have preserved all valid topic partition to consumer assignments and removed
              // all invalid topic partitions and invalid consumers. Now we need to assign unassignedPartitions
              // to consumers so that the topic partition assignments are as balanced as possible.
              sortedCurrentSubscriptions.addAll(currentAssignment.keySet());
  
              balance(prevAssignment, unassignedPartitions);
  
              log.info("Final assignment of partitions to consumers: \n{}", currentAssignment);
  
              return currentAssignment;
          }
  
          private List<TopicPartition> assignOwnedPartitions() {
              List<TopicPartition> assignedPartitions = new ArrayList<>();
              for (Iterator<Entry<String, List<TopicPartition>>> it = currentAssignment.entrySet().iterator(); it.hasNext();) {
                  Map.Entry<String, List<TopicPartition>> entry = it.next();
                  String consumer = entry.getKey();
                  Subscription consumerSubscription = subscriptions.get(consumer);
                  if (consumerSubscription == null) {
                      // if a consumer that existed before (and had some partition assignments) is now removed, remove it from currentAssignment
                      for (TopicPartition topicPartition: entry.getValue())
                          currentPartitionConsumer.remove(topicPartition);
                      it.remove();
                  } else {
                      // otherwise (the consumer still exists)
                      for (Iterator<TopicPartition> partitionIter = entry.getValue().iterator(); partitionIter.hasNext();) {
                          TopicPartition partition = partitionIter.next();
                          if (!topic2AllPotentialConsumers.containsKey(partition.topic())) {
                              // if this topic partition of this consumer no longer exists, remove it from currentAssignment of the consumer
                              partitionIter.remove();
                              currentPartitionConsumer.remove(partition);
                          } else if (!consumerSubscription.topics().contains(partition.topic()) || rackInfo.racksMismatch(consumer, partition)) {
                              // if the consumer is no longer subscribed to its topic or if racks don't match for rack-aware assignment,
                              // remove it from currentAssignment of the consumer
                              partitionIter.remove();
                              revocationRequired = true;
                          } else {
                              // otherwise, remove the topic partition from those that need to be assigned only if
                              // its current consumer is still subscribed to its topic (because it is already assigned
                              // and we would want to preserve that assignment as much as possible)
                              assignedPartitions.add(partition);
                          }
                      }
                  }
              }
              return assignedPartitions;
          }
  
          /**
           * get the unassigned partition list by computing the difference set of the sortedPartitions(all partitions)
           * and sortedAssignedPartitions. If no assigned partitions, we'll just return all sorted topic partitions.
           *
           * We loop the sortedPartition, and compare the ith element in sortedAssignedPartitions(i start from 0):
           *   - if not equal to the ith element, add to unassignedPartitions
           *   - if equal to the ith element, get next element from sortedAssignedPartitions
           *
           * @param sortedAssignedPartitions:     sorted partitions, all are included in the sortedPartitions
           * @return                              partitions that aren't assigned to any current consumer
           */
          private List<TopicPartition> getUnassignedPartitions(List<TopicPartition> sortedAssignedPartitions) {
              if (sortedAssignedPartitions.isEmpty()) {
                  return sortedAllPartitions;
              }
  
              List<TopicPartition> unassignedPartitions = new ArrayList<>();
  
              sortedAssignedPartitions.sort(new PartitionComparator(topic2AllPotentialConsumers));
  
              boolean shouldAddDirectly = false;
              Iterator<TopicPartition> sortedAssignedPartitionsIter = sortedAssignedPartitions.iterator();
              TopicPartition nextAssignedPartition = sortedAssignedPartitionsIter.next();
  
              for (TopicPartition topicPartition : sortedAllPartitions) {
                  if (shouldAddDirectly || !nextAssignedPartition.equals(topicPartition)) {
                      unassignedPartitions.add(topicPartition);
                  } else {
                      // this partition is in assignedPartitions, don't add to unassignedPartitions, just get next assigned partition
                      if (sortedAssignedPartitionsIter.hasNext()) {
                          nextAssignedPartition = sortedAssignedPartitionsIter.next();
                      } else {
                          // add the remaining directly since there is no more sortedAssignedPartitions
                          shouldAddDirectly = true;
                      }
                  }
              }
              return unassignedPartitions;
          }
  
          /**
           * update the prevAssignment with the partitions, consumer and generation in parameters
           *
           * @param partitions:       The partitions to be updated the prevAssignment
           * @param consumer:         The consumer Id
           * @param prevAssignment:   The assignment contains the assignment with the 2nd largest generation
           * @param generation:       The generation of this assignment (partitions)
           */
          private void updatePrevAssignment(Map<TopicPartition, ConsumerGenerationPair> prevAssignment,
                                            List<TopicPartition> partitions,
                                            String consumer,
                                            int generation) {
              for (TopicPartition partition: partitions) {
                  if (prevAssignment.containsKey(partition)) {
                      // only keep the latest previous assignment
                      if (generation > prevAssignment.get(partition).generation) {
                          prevAssignment.put(partition, new ConsumerGenerationPair(consumer, generation));
                      }
                  } else {
                      prevAssignment.put(partition, new ConsumerGenerationPair(consumer, generation));
                  }
              }
          }
  
          /**
           * filling in the prevAssignment from the subscriptions.
           *
           * @param prevAssignment: The assignment contains the assignment with the 2nd largest generation
           */
          private void prepopulateCurrentAssignments(Map<TopicPartition, ConsumerGenerationPair> prevAssignment) {
              // we need to process subscriptions' user data with each consumer's reported generation in mind
              // higher generations overwrite lower generations in case of a conflict
              // note that a conflict could exist only if user data is for different generations
  
              for (Map.Entry<String, Subscription> subscriptionEntry: subscriptions.entrySet()) {
                  String consumer = subscriptionEntry.getKey();
                  Subscription subscription = subscriptionEntry.getValue();
                  if (subscription.userData() != null) {
                      // since this is our 2nd time to deserialize memberData, rewind userData is necessary
                      subscription.userData().rewind();
                  }
  
                  MemberData memberData = memberData(subscription);
  
                  // we already have the maxGeneration info, so just compare the current generation of memberData, and put into prevAssignment
                  if (memberData.generation.isPresent() && memberData.generation.get() < maxGeneration) {
                      // if the current member's generation is lower than maxGeneration, put into prevAssignment if needed
                      updatePrevAssignment(prevAssignment, memberData.partitions, consumer, memberData.generation.get());
                  } else if (memberData.generation.isEmpty() && maxGeneration > DEFAULT_GENERATION) {
                      // if maxGeneration is larger than DEFAULT_GENERATION
                      // put all (no generation) partitions as DEFAULT_GENERATION into prevAssignment if needed
                      updatePrevAssignment(prevAssignment, memberData.partitions, consumer, DEFAULT_GENERATION);
                  }
              }
          }
  
          /**
           * determine if the current assignment is a balanced one
           *
           * @return true if the given assignment is balanced; false otherwise
           */
          private boolean isBalanced() {
              int min = currentAssignment.get(sortedCurrentSubscriptions.first()).size();
              int max = currentAssignment.get(sortedCurrentSubscriptions.last()).size();
              if (min >= max - 1)
                  // if minimum and maximum numbers of partitions assigned to consumers differ by at most one return true
                  return true;
  
              // create a mapping from partitions to the consumer assigned to them
              final Map<TopicPartition, String> allPartitions = new HashMap<>();
              Set<Entry<String, List<TopicPartition>>> assignments = currentAssignment.entrySet();
              for (Map.Entry<String, List<TopicPartition>> entry: assignments) {
                  List<TopicPartition> topicPartitions = entry.getValue();
                  for (TopicPartition topicPartition: topicPartitions) {
                      if (allPartitions.containsKey(topicPartition))
                          log.error("{} is assigned to more than one consumer.", topicPartition);
                      allPartitions.put(topicPartition, entry.getKey());
                  }
              }
  
              // for each consumer that does not have all the topic partitions it can get make sure none of the topic partitions it
              // could but did not get cannot be moved to it (because that would break the balance)
              for (String consumer: sortedCurrentSubscriptions) {
                  List<TopicPartition> consumerPartitions = currentAssignment.get(consumer);
                  int consumerPartitionCount = consumerPartitions.size();
  
                  // skip if this consumer already has all the topic partitions it can get
                  List<String> allSubscribedTopics = consumer2AllPotentialTopics.get(consumer);
                  int maxAssignmentSize = getMaxAssignmentSize(allSubscribedTopics);
  
                  if (consumerPartitionCount == maxAssignmentSize)
                      continue;
  
                  // otherwise make sure it cannot get any more
                  for (String topic: allSubscribedTopics) {
                      int partitionCount = partitionsPerTopic.get(topic).size();
                      for (int i = 0; i < partitionCount; i++) {
                          TopicPartition topicPartition = new TopicPartition(topic, i);
                          if (!currentAssignment.get(consumer).contains(topicPartition)) {
                              String otherConsumer = allPartitions.get(topicPartition);
                              int otherConsumerPartitionCount = currentAssignment.get(otherConsumer).size();
                              if (consumerPartitionCount + 1 < otherConsumerPartitionCount) {
                                  log.debug("{} can be moved from consumer {} to consumer {} for a more balanced assignment.",
                                          topicPartition, otherConsumer, consumer);
                                  return false;
                              }
                          }
                      }
                  }
              }
              return true;
          }
  
          /**
           * get the maximum assigned partition size of the {@code allSubscribedTopics}
           *
           * @param allSubscribedTopics           the subscribed topics of a consumer
           * @return                              maximum assigned partition size
           */
          private int getMaxAssignmentSize(List<String> allSubscribedTopics) {
              int maxAssignmentSize;
              if (allSubscribedTopics.size() == partitionsPerTopic.size()) {
                  maxAssignmentSize = totalPartitionsCount;
              } else {
                  maxAssignmentSize = allSubscribedTopics.stream().map(partitionsPerTopic::get).map(List::size).reduce(0, Integer::sum);
              }
              return maxAssignmentSize;
          }
  
          /**
           * @return the balance score of the given assignment, as the sum of assigned partitions size difference of all consumer pairs.
           * A perfectly balanced assignment (with all consumers getting the same number of partitions) has a balance score of 0.
           * Lower balance score indicates a more balanced assignment.
           */
          private int getBalanceScore(Map<String, List<TopicPartition>> assignment) {
              int score = 0;
  
              Map<String, Integer> consumer2AssignmentSize = new HashMap<>();
              for (Entry<String, List<TopicPartition>> entry: assignment.entrySet())
                  consumer2AssignmentSize.put(entry.getKey(), entry.getValue().size());
  
              Iterator<Entry<String, Integer>> it = consumer2AssignmentSize.entrySet().iterator();
              while (it.hasNext()) {
                  Entry<String, Integer> entry = it.next();
                  int consumerAssignmentSize = entry.getValue();
                  it.remove();
                  for (Entry<String, Integer> otherEntry: consumer2AssignmentSize.entrySet())
                      score += Math.abs(consumerAssignmentSize - otherEntry.getValue());
              }
  
              return score;
          }
  
          /**
           * The assignment should improve the overall balance of the partition assignments to consumers.
           */
          private boolean maybeAssignPartition(TopicPartition partition, RackInfo rackInfo) {
              for (String consumer: sortedCurrentSubscriptions) {
                  if (consumer2AllPotentialTopics.get(consumer).contains(partition.topic()) && (rackInfo == null || !rackInfo.racksMismatch(consumer, partition))) {
                      sortedCurrentSubscriptions.remove(consumer);
                      currentAssignment.get(consumer).add(partition);
                      currentPartitionConsumer.put(partition, consumer);
                      sortedCurrentSubscriptions.add(consumer);
                      return true;
                  }
              }
              return false;
          }
  
          /**
           * attempt to assign all unassigned partitions
           *
           * @param unassignedPartitions partitions that are still unassigned
           * @param rackInfo             rack information used to match racks. If null, no rack-matching is performed
           * @param removeAssigned       flag that indicates if assigned partitions should be removed from `unassignedPartitions`
           */
          private void maybeAssign(List<TopicPartition> unassignedPartitions, RackInfo rackInfo, boolean removeAssigned) {
              // assign all unassigned partitions
              for (Iterator<TopicPartition> iter = unassignedPartitions.iterator(); iter.hasNext();) {
                  TopicPartition partition = iter.next();
                  // skip if there is no potential consumer for the topic
                  if (topic2AllPotentialConsumers.get(partition.topic()).isEmpty())
                      continue;
  
                  if (maybeAssignPartition(partition, rackInfo) && removeAssigned)
                      iter.remove();
              }
          }
  
          private boolean canTopicParticipateInReassignment(String topic) {
              // if a topic has two or more potential consumers it is subject to reassignment.
              return topic2AllPotentialConsumers.get(topic).size() >= 2;
          }
  
          private boolean canConsumerParticipateInReassignment(String consumer) {
              List<TopicPartition> currentPartitions = currentAssignment.get(consumer);
              int currentAssignmentSize = currentPartitions.size();
              List<String> allSubscribedTopics = consumer2AllPotentialTopics.get(consumer);
              int maxAssignmentSize = getMaxAssignmentSize(allSubscribedTopics);
  
              if (currentAssignmentSize > maxAssignmentSize)
                  log.error("The consumer {} is assigned more partitions than the maximum possible.", consumer);
  
              if (currentAssignmentSize < maxAssignmentSize)
                  // if a consumer is not assigned all its potential partitions it is subject to reassignment
                  return true;
  
              for (TopicPartition partition: currentPartitions)
                  // if any of the partitions assigned to a consumer is subject to reassignment the consumer itself
                  // is subject to reassignment
                  if (canTopicParticipateInReassignment(partition.topic()))
                      return true;
  
              return false;
          }
  
          /**
           * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
           */
          private void balance(Map<TopicPartition, ConsumerGenerationPair> prevAssignment,
                               List<TopicPartition> unassignedPartitions) {
              boolean initializing = currentAssignment.get(sortedCurrentSubscriptions.last()).isEmpty();
  
              // First assign with rack matching and then assign any remaining without rack matching
              List<TopicPartition> partitionsToAssign = unassignedPartitions;
              if (!rackInfo.consumerRacks.isEmpty()) {
                  partitionsToAssign = new LinkedList<>(unassignedPartitions);
                  maybeAssign(partitionsToAssign, rackInfo, true);
              }
              maybeAssign(partitionsToAssign, null, false);
  
              // narrow down the reassignment scope to only those partitions that can actually be reassigned
              Set<TopicPartition> fixedPartitions = new HashSet<>();
              for (String topic: topic2AllPotentialConsumers.keySet())
                  if (!canTopicParticipateInReassignment(topic)) {
                      for (int i = 0; i < partitionsPerTopic.get(topic).size(); i++) {
                          fixedPartitions.add(new TopicPartition(topic, i));
                      }
                  }
              sortedAllPartitions.removeAll(fixedPartitions);
              unassignedPartitions.removeAll(fixedPartitions);
  
              // narrow down the reassignment scope to only those consumers that are subject to reassignment
              Map<String, List<TopicPartition>> fixedAssignments = new HashMap<>();
              for (String consumer: consumer2AllPotentialTopics.keySet())
                  if (!canConsumerParticipateInReassignment(consumer)) {
                      sortedCurrentSubscriptions.remove(consumer);
                      fixedAssignments.put(consumer, currentAssignment.remove(consumer));
                  }
  
              // create a deep copy of the current assignment so we can revert to it if we do not get a more balanced assignment later
              Map<String, List<TopicPartition>> preBalanceAssignment = deepCopy(currentAssignment);
              Map<TopicPartition, String> preBalancePartitionConsumers = new HashMap<>(currentPartitionConsumer);
  
              // if we don't already need to revoke something due to subscription changes, first try to balance by only moving newly added partitions
              if (!revocationRequired) {
                  performReassignments(unassignedPartitions, prevAssignment);
              }
  
              boolean reassignmentPerformed = performReassignments(sortedAllPartitions, prevAssignment);
  
              // if we are not preserving existing assignments and we have made changes to the current assignment
              // make sure we are getting a more balanced assignment; otherwise, revert to previous assignment
              if (!initializing && reassignmentPerformed && getBalanceScore(currentAssignment) >= getBalanceScore(preBalanceAssignment)) {
                  deepCopy(preBalanceAssignment, currentAssignment);
                  currentPartitionConsumer.clear();
                  currentPartitionConsumer.putAll(preBalancePartitionConsumers);
              }
  
              // add the fixed assignments (those that could not change) back
              for (Entry<String, List<TopicPartition>> entry: fixedAssignments.entrySet()) {
                  String consumer = entry.getKey();
                  currentAssignment.put(consumer, entry.getValue());
                  sortedCurrentSubscriptions.add(consumer);
              }
  
              fixedAssignments.clear();
          }
  
          private boolean performReassignments(List<TopicPartition> reassignablePartitions,
                                               Map<TopicPartition, ConsumerGenerationPair> prevAssignment) {
              boolean reassignmentPerformed = false;
              boolean modified;
  
              // repeat reassignment until no partition can be moved to improve the balance
              do {
                  modified = false;
                  // reassign all reassignable partitions (starting from the partition with least potential consumers and if needed)
                  // until the full list is processed or a balance is achieved
                  Iterator<TopicPartition> partitionIterator = reassignablePartitions.iterator();
                  while (partitionIterator.hasNext() && !isBalanced()) {
                      TopicPartition partition = partitionIterator.next();
  
                      // the partition must have at least two consumers
                      if (topic2AllPotentialConsumers.get(partition.topic()).size() <= 1)
                          log.error("Expected more than one potential consumer for partition '{}'", partition);
  
                      // the partition must have a current consumer
                      String consumer = currentPartitionConsumer.get(partition);
                      if (consumer == null)
                          log.error("Expected partition '{}' to be assigned to a consumer", partition);
  
                      if (prevAssignment.containsKey(partition) &&
                              currentAssignment.get(consumer).size() > currentAssignment.get(prevAssignment.get(partition).consumer).size() + 1) {
                          reassignPartition(partition, prevAssignment.get(partition).consumer);
                          reassignmentPerformed = true;
                          modified = true;
                          continue;
                      }
  
                      // check if a better-suited consumer exist for the partition; if so, reassign it
                      // Use consumer within rack if possible
                      String consumerRack = rackInfo.consumerRacks.get(consumer);
                      Set<String> partitionRacks = rackInfo.partitionRacks.get(partition);
                      boolean foundRackConsumer = false;
                      if (consumerRack != null && !partitionRacks.isEmpty() && partitionRacks.contains(consumerRack)) {
                          for (String otherConsumer : topic2AllPotentialConsumers.get(partition.topic())) {
                              String otherConsumerRack = rackInfo.consumerRacks.get(otherConsumer);
                              if (otherConsumerRack == null || !partitionRacks.contains(otherConsumerRack))
                                  continue;
                              if (currentAssignment.get(consumer).size() > currentAssignment.get(otherConsumer).size() + 1) {
                                  reassignPartition(partition);
                                  reassignmentPerformed = true;
                                  modified = true;
                                  foundRackConsumer = true;
                                  break;
                              }
                          }
                      }
                      if (!foundRackConsumer) {
                          for (String otherConsumer : topic2AllPotentialConsumers.get(partition.topic())) {
                              if (currentAssignment.get(consumer).size() > currentAssignment.get(otherConsumer).size() + 1) {
                                  reassignPartition(partition);
                                  reassignmentPerformed = true;
                                  modified = true;
                                  break;
                              }
                          }
                      }
                  }
              } while (modified);
  
              return reassignmentPerformed;
          }
  
          private void reassignPartition(TopicPartition partition) {
              // find the new consumer
              String newConsumer = null;
              for (String anotherConsumer: sortedCurrentSubscriptions) {
                  if (consumer2AllPotentialTopics.get(anotherConsumer).contains(partition.topic())) {
                      newConsumer = anotherConsumer;
                      break;
                  }
              }
  
              assert newConsumer != null;
  
              reassignPartition(partition, newConsumer);
          }
  
          private void reassignPartition(TopicPartition partition, String newConsumer) {
              String consumer = currentPartitionConsumer.get(partition);
              // find the correct partition movement considering the stickiness requirement
              TopicPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(partition, consumer, newConsumer);
              processPartitionMovement(partitionToBeMoved, newConsumer);
          }
  
          private void processPartitionMovement(TopicPartition partition, String newConsumer) {
              String oldConsumer = currentPartitionConsumer.get(partition);
  
              sortedCurrentSubscriptions.remove(oldConsumer);
              sortedCurrentSubscriptions.remove(newConsumer);
  
              partitionMovements.movePartition(partition, oldConsumer, newConsumer);
  
              currentAssignment.get(oldConsumer).remove(partition);
              currentAssignment.get(newConsumer).add(partition);
              currentPartitionConsumer.put(partition, newConsumer);
              sortedCurrentSubscriptions.add(newConsumer);
              sortedCurrentSubscriptions.add(oldConsumer);
          }
  
          private void deepCopy(Map<String, List<TopicPartition>> source, Map<String, List<TopicPartition>> dest) {
              dest.clear();
              for (Entry<String, List<TopicPartition>> entry: source.entrySet())
                  dest.put(entry.getKey(), new ArrayList<>(entry.getValue()));
          }
  
          private Map<String, List<TopicPartition>> deepCopy(Map<String, List<TopicPartition>> assignment) {
              Map<String, List<TopicPartition>> copy = new HashMap<>();
              deepCopy(assignment, copy);
              return copy;
          }
      }
  }
  