diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java index 5798909927461..e111aa62599e4 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java @@ -29,7 +29,6 @@ import java.util.Map.Entry; import java.util.Optional; import java.util.Set; -import java.util.TreeMap; import java.util.TreeSet; import java.util.stream.Collectors; import org.apache.kafka.common.TopicPartition; @@ -40,6 +39,7 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor { private static final Logger log = LoggerFactory.getLogger(AbstractStickyAssignor.class); public static final int DEFAULT_GENERATION = -1; + public int maxGeneration = DEFAULT_GENERATION; private PartitionMovements partitionMovements; @@ -77,12 +77,10 @@ public Map> assign(Map partitionsP partitionsTransferringOwnership = new HashMap<>(); return constrainedAssign(partitionsPerTopic, consumerToOwnedPartitions); } else { - log.debug("Detected that all not consumers were subscribed to same set of topics, falling back to the " + log.debug("Detected that not all consumers were subscribed to same set of topics, falling back to the " + "general case assignment algorithm"); partitionsTransferringOwnership = null; - // we don't need consumerToOwnedPartitions in general assign case - consumerToOwnedPartitions = null; - return generalAssign(partitionsPerTopic, subscriptions); + return generalAssign(partitionsPerTopic, subscriptions, consumerToOwnedPartitions); } } @@ -95,7 +93,7 @@ private boolean allSubscriptionsEqual(Set allTopics, Map> consumerToOwnedPartitions) { Set membersWithOldGeneration = new HashSet<>(); Set membersOfCurrentHighestGeneration = new HashSet<>(); - int maxGeneration = DEFAULT_GENERATION; + boolean isAllSubscriptionsEqual = true; Set subscribedTopics = new HashSet<>(); @@ -106,9 +104,9 @@ private boolean allSubscriptionsEqual(Set allTopics, // initialize the subscribed topics set if this is the first subscription if (subscribedTopics.isEmpty()) { subscribedTopics.addAll(subscription.topics()); - } else if (!(subscription.topics().size() == subscribedTopics.size() + } else if (isAllSubscriptionsEqual && !(subscription.topics().size() == subscribedTopics.size() && subscribedTopics.containsAll(subscription.topics()))) { - return false; + isAllSubscriptionsEqual = false; } MemberData memberData = memberData(subscription); @@ -141,7 +139,7 @@ private boolean allSubscriptionsEqual(Set allTopics, for (String consumer : membersWithOldGeneration) { consumerToOwnedPartitions.get(consumer).clear(); } - return true; + return isAllSubscriptionsEqual; } @@ -301,62 +299,6 @@ private Map> constrainedAssign(Map return assignment; } - /** - * 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 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 the ith element, get next element from sortedAssignedPartitions - * - * @param totalPartitionsCount all partitions counts in this assignment - * @param partitionsPerTopic the number of partitions for each subscribed topic. - * @param sortedAssignedPartitions sorted partitions, all are included in the sortedPartitions - * @return the partitions not yet assigned to any consumers - */ - private List getUnassignedPartitions(int totalPartitionsCount, - Map partitionsPerTopic, - List sortedAssignedPartitions) { - List 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(partitionsPerTopic, sortedAllTopics, totalPartitionsCount); - } - - List unassignedPartitions = new ArrayList<>(totalPartitionsCount - sortedAssignedPartitions.size()); - - Collections.sort(sortedAssignedPartitions, Comparator.comparing(TopicPartition::topic).thenComparing(TopicPartition::partition)); - - boolean shouldAddDirectly = false; - Iterator sortedAssignedPartitionsIter = sortedAssignedPartitions.iterator(); - TopicPartition nextAssignedPartition = sortedAssignedPartitionsIter.next(); - - for (String topic : sortedAllTopics) { - int partitionCount = partitionsPerTopic.get(topic); - 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; - } - private List getAllTopicPartitions(Map partitionsPerTopic, List sortedAllTopics, @@ -384,37 +326,39 @@ private List getAllTopicPartitions(Map partitio * * @param partitionsPerTopic The number of 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 * * @return Map from each member to the list of partitions assigned to them. */ private Map> generalAssign(Map partitionsPerTopic, - Map subscriptions) { - Map> currentAssignment = new HashMap<>(); + Map subscriptions, + Map> currentAssignment) { + if (log.isDebugEnabled()) { + log.debug("performing general assign. partitionsPerTopic: {}, subscriptions: {}, currentAssignment: {}", + partitionsPerTopic, subscriptions, currentAssignment); + } + Map prevAssignment = new HashMap<>(); partitionMovements = new PartitionMovements(); - prepopulateCurrentAssignments(subscriptions, currentAssignment, prevAssignment); + prepopulateCurrentAssignments(subscriptions, prevAssignment); - // a mapping of all topic partitions to all consumers that can be assigned to them - final Map> partition2AllPotentialConsumers = new HashMap<>(); - // a mapping of all consumers to all potential topic partitions that can be assigned to them - final Map> consumer2AllPotentialPartitions = new HashMap<>(); + // a mapping of all topics to all consumers that can be assigned to them + final Map> topic2AllPotentialConsumers = new HashMap<>(partitionsPerTopic.keySet().size()); + // a mapping of all consumers to all potential topics that can be assigned to them + final Map> consumer2AllPotentialTopics = new HashMap<>(subscriptions.keySet().size()); - // initialize partition2AllPotentialConsumers and consumer2AllPotentialPartitions in the following two for loops - for (Entry entry: partitionsPerTopic.entrySet()) { - for (int i = 0; i < entry.getValue(); ++i) - partition2AllPotentialConsumers.put(new TopicPartition(entry.getKey(), i), new ArrayList<>()); - } + // initialize topic2AllPotentialConsumers and consumer2AllPotentialTopics + partitionsPerTopic.keySet().stream().forEach( + topicName -> topic2AllPotentialConsumers.put(topicName, new ArrayList<>())); for (Entry entry: subscriptions.entrySet()) { String consumerId = entry.getKey(); - consumer2AllPotentialPartitions.put(consumerId, new ArrayList<>()); + List subscribedTopics = new ArrayList<>(entry.getValue().topics().size()); + consumer2AllPotentialTopics.put(consumerId, subscribedTopics); entry.getValue().topics().stream().filter(topic -> partitionsPerTopic.get(topic) != null).forEach(topic -> { - for (int i = 0; i < partitionsPerTopic.get(topic); ++i) { - TopicPartition topicPartition = new TopicPartition(topic, i); - consumer2AllPotentialPartitions.get(consumerId).add(topicPartition); - partition2AllPotentialConsumers.get(topicPartition).add(consumerId); - } + 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 @@ -428,14 +372,18 @@ private Map> generalAssign(Map par for (TopicPartition topicPartition: entry.getValue()) currentPartitionConsumer.put(topicPartition, entry.getKey()); - List sortedPartitions = sortPartitions(partition2AllPotentialConsumers); + int totalPartitionsCount = partitionsPerTopic.values().stream().reduce(0, Integer::sum); + List sortedAllTopics = new ArrayList<>(topic2AllPotentialConsumers.keySet()); + Collections.sort(sortedAllTopics, new TopicComparator(topic2AllPotentialConsumers)); + List sortedAllPartitions = getAllTopicPartitions(partitionsPerTopic, sortedAllTopics, totalPartitionsCount); - // all partitions that need to be assigned (initially set to all partitions but adjusted in the following loop) - List unassignedPartitions = new ArrayList<>(sortedPartitions); + // the partitions already assigned in current assignment + List assignedPartitions = new ArrayList<>(); boolean revocationRequired = false; for (Iterator>> it = currentAssignment.entrySet().iterator(); it.hasNext();) { Map.Entry> entry = it.next(); - if (!subscriptions.containsKey(entry.getKey())) { + Subscription consumerSubscription = subscriptions.get(entry.getKey()); + 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); @@ -444,23 +392,32 @@ private Map> generalAssign(Map par // otherwise (the consumer still exists) for (Iterator partitionIter = entry.getValue().iterator(); partitionIter.hasNext();) { TopicPartition partition = partitionIter.next(); - if (!partition2AllPotentialConsumers.containsKey(partition)) { - // if this topic partition of this consumer no longer exists remove it from currentAssignment of the consumer + 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 (!subscriptions.get(entry.getKey()).topics().contains(partition.topic())) { - // if this partition cannot remain assigned to its current consumer because the consumer - // is no longer subscribed to its topic remove it from currentAssignment of the consumer + } else if (!consumerSubscription.topics().contains(partition.topic())) { + // because the consumer is no longer subscribed to its topic, remove it from currentAssignment of the consumer partitionIter.remove(); revocationRequired = true; - } else + } 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) - unassignedPartitions.remove(partition); + assignedPartitions.add(partition); + } } } } + + // all partitions that needed to be assigned + List unassignedPartitions = getUnassignedPartitions(sortedAllPartitions, assignedPartitions, topic2AllPotentialConsumers); + assignedPartitions = null; + + 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. @@ -469,58 +426,172 @@ private Map> generalAssign(Map par TreeSet sortedCurrentSubscriptions = new TreeSet<>(new SubscriptionComparator(currentAssignment)); sortedCurrentSubscriptions.addAll(currentAssignment.keySet()); - balance(currentAssignment, prevAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions, - consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer, revocationRequired); + balance(currentAssignment, prevAssignment, sortedAllPartitions, unassignedPartitions, sortedCurrentSubscriptions, + consumer2AllPotentialTopics, topic2AllPotentialConsumers, currentPartitionConsumer, revocationRequired, + partitionsPerTopic, totalPartitionsCount); + + if (log.isDebugEnabled()) { + log.debug("final assignment: {}", currentAssignment); + } + return currentAssignment; } + /** + * 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. + * This is used in generalAssign method + * + * 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 the ith element, get next element from sortedAssignedPartitions + * + * @param sortedAllPartitions: sorted all partitions + * @param sortedAssignedPartitions: sorted partitions, all are included in the sortedPartitions + * @param topic2AllPotentialConsumers: topics mapped to all consumers that subscribed to it + * @return partitions that aren't assigned to any current consumer + */ + private List getUnassignedPartitions(List sortedAllPartitions, + List sortedAssignedPartitions, + Map> topic2AllPotentialConsumers) { + if (sortedAssignedPartitions.isEmpty()) { + return sortedAllPartitions; + } + + List unassignedPartitions = new ArrayList<>(); + + Collections.sort(sortedAssignedPartitions, new PartitionComparator(topic2AllPotentialConsumers)); + + boolean shouldAddDirectly = false; + Iterator 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; + } + + /** + * 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. + * This is used in constrainedAssign method + * + * 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 the ith element, get next element from sortedAssignedPartitions + * + * @param totalPartitionsCount all partitions counts in this assignment + * @param partitionsPerTopic the number of partitions for each subscribed topic. + * @param sortedAssignedPartitions sorted partitions, all are included in the sortedPartitions + * @return the partitions not yet assigned to any consumers + */ + private List getUnassignedPartitions(int totalPartitionsCount, + Map partitionsPerTopic, + List sortedAssignedPartitions) { + List 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(partitionsPerTopic, sortedAllTopics, totalPartitionsCount); + } + + List unassignedPartitions = new ArrayList<>(totalPartitionsCount - sortedAssignedPartitions.size()); + + Collections.sort(sortedAssignedPartitions, Comparator.comparing(TopicPartition::topic).thenComparing(TopicPartition::partition)); + + boolean shouldAddDirectly = false; + Iterator sortedAssignedPartitionsIter = sortedAssignedPartitions.iterator(); + TopicPartition nextAssignedPartition = sortedAssignedPartitionsIter.next(); + + for (String topic : sortedAllTopics) { + int partitionCount = partitionsPerTopic.get(topic); + 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; + } + + /** + * update the prevAssignment with the partitions, consumer and generation in parameters + * + * @param partitions: The partitions to be updated the prevAssignement + * @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 prevAssignment, + List 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 subscriptions: Map from the member id to their respective topic subscription + * @param prevAssignment: The assignment contains the assignment with the 2nd largest generation + */ private void prepopulateCurrentAssignments(Map subscriptions, - Map> currentAssignment, Map 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 exists only if user data is for different generations - // for each partition we create a sorted map of its consumers by generation - Map> sortedPartitionConsumersByGeneration = new HashMap<>(); for (Map.Entry subscriptionEntry: subscriptions.entrySet()) { String consumer = subscriptionEntry.getKey(); - MemberData memberData = memberData(subscriptionEntry.getValue()); - - for (TopicPartition partition: memberData.partitions) { - if (sortedPartitionConsumersByGeneration.containsKey(partition)) { - Map consumers = sortedPartitionConsumersByGeneration.get(partition); - if (memberData.generation.isPresent() && consumers.containsKey(memberData.generation.get())) { - // same partition is assigned to two consumers during the same rebalance. - // log a warning and skip this record - log.warn("Partition '{}' is assigned to multiple consumers following sticky assignment generation {}.", - partition, memberData.generation); - } else - consumers.put(memberData.generation.orElse(DEFAULT_GENERATION), consumer); - } else { - TreeMap sortedConsumers = new TreeMap<>(); - sortedConsumers.put(memberData.generation.orElse(DEFAULT_GENERATION), consumer); - sortedPartitionConsumersByGeneration.put(partition, sortedConsumers); - } + 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(subscriptionEntry.getValue()); - // prevAssignment holds the prior ConsumerGenerationPair (before current) of each partition - // current and previous consumers are the last two consumers of each partition in the above sorted map - for (Map.Entry> partitionConsumersEntry: sortedPartitionConsumersByGeneration.entrySet()) { - TopicPartition partition = partitionConsumersEntry.getKey(); - TreeMap consumers = partitionConsumersEntry.getValue(); - Iterator it = consumers.descendingKeySet().iterator(); - - // let's process the current (most recent) consumer first - String consumer = consumers.get(it.next()); - currentAssignment.computeIfAbsent(consumer, k -> new ArrayList<>()); - currentAssignment.get(consumer).add(partition); - - // now update previous assignment if any - if (it.hasNext()) { - int generation = it.next(); - prevAssignment.put(partition, new ConsumerGenerationPair(consumers.get(generation), generation)); + // 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.isPresent() && maxGeneration > DEFAULT_GENERATION) { + // if maxGeneration is larger then DEFAULT_GENERATION + // put all (no generation) partitions as DEFAULT_GENERATION into prevAssignment if needed + updatePrevAssignment(prevAssignment, memberData.partitions, consumer, DEFAULT_GENERATION); } } } @@ -528,14 +599,18 @@ private void prepopulateCurrentAssignments(Map subscriptio /** * determine if the current assignment is a balanced one * - * @param currentAssignment the assignment whose balance needs to be checked - * @param sortedCurrentSubscriptions an ascending sorted set of consumers based on how many topic partitions are already assigned to them - * @param allSubscriptions a mapping of all consumers to all potential topic partitions that can be assigned to them + * @param currentAssignment: the assignment whose balance needs to be checked + * @param sortedCurrentSubscriptions: an ascending sorted set of consumers based on how many topic partitions are already assigned to them + * @param allSubscriptions: a mapping of all consumers to all potential topics that can be assigned to them + * @param partitionsPerTopic: The number of partitions for each subscribed topic + * @param totalPartitionCount total partition count to be assigned * @return true if the given assignment is balanced; false otherwise */ private boolean isBalanced(Map> currentAssignment, TreeSet sortedCurrentSubscriptions, - Map> allSubscriptions) { + Map> allSubscriptions, + Map partitionsPerTopic, + int totalPartitionCount) { int min = currentAssignment.get(sortedCurrentSubscriptions.first()).size(); int max = currentAssignment.get(sortedCurrentSubscriptions.last()).size(); if (min >= max - 1) @@ -561,19 +636,25 @@ private boolean isBalanced(Map> currentAssignment, int consumerPartitionCount = consumerPartitions.size(); // skip if this consumer already has all the topic partitions it can get - if (consumerPartitionCount == allSubscriptions.get(consumer).size()) + List allSubscribedTopics = allSubscriptions.get(consumer); + int maxAssignmentSize = getMaxAssignmentSize(totalPartitionCount, allSubscribedTopics, partitionsPerTopic); + + if (consumerPartitionCount == maxAssignmentSize) continue; // otherwise make sure it cannot get any more - List potentialTopicPartitions = allSubscriptions.get(consumer); - for (TopicPartition topicPartition: potentialTopicPartitions) { - if (!currentAssignment.get(consumer).contains(topicPartition)) { - String otherConsumer = allPartitions.get(topicPartition); - int otherConsumerPartitionCount = currentAssignment.get(otherConsumer).size(); - if (consumerPartitionCount < otherConsumerPartitionCount) { - log.debug("{} can be moved from consumer {} to consumer {} for a more balanced assignment.", - topicPartition, otherConsumer, consumer); - return false; + for (String topic: allSubscribedTopics) { + int partitionCount = partitionsPerTopic.get(topic); + 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 < otherConsumerPartitionCount) { + log.debug("{} can be moved from consumer {} to consumer {} for a more balanced assignment.", + topicPartition, otherConsumer, consumer); + return false; + } } } } @@ -581,6 +662,26 @@ private boolean isBalanced(Map> currentAssignment, return true; } + /** + * get the maximum assigned partition size of the {@code allSubscribedTopics} + * + * @param totalPartitionCount total partition count to be assigned + * @param allSubscribedTopics the subscribed topics of a consumer + * @param partitionsPerTopic The number of partitions for each subscribed topic + * @return maximum assigned partition size + */ + private int getMaxAssignmentSize(int totalPartitionCount, + List allSubscribedTopics, + Map partitionsPerTopic) { + int maxAssignmentSize; + if (allSubscribedTopics.size() == partitionsPerTopic.size()) { + maxAssignmentSize = totalPartitionCount; + } else { + maxAssignmentSize = allSubscribedTopics.stream().map(topic -> partitionsPerTopic.get(topic)).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. @@ -605,29 +706,16 @@ private int getBalanceScore(Map> assignment) { return score; } - /** - * Sort valid partitions so they are processed in the potential reassignment phase in the proper order - * that causes minimal partition movement among consumers (hence honoring maximal stickiness) - * - * @param partition2AllPotentialConsumers a mapping of partitions to their potential consumers - * @return an ascending sorted list of topic partitions based on how many consumers can potentially use them - */ - private List sortPartitions(Map> partition2AllPotentialConsumers) { - List sortedPartitions = new ArrayList<>(partition2AllPotentialConsumers.keySet()); - Collections.sort(sortedPartitions, new PartitionComparator(partition2AllPotentialConsumers)); - return sortedPartitions; - } - /** * The assignment should improve the overall balance of the partition assignments to consumers. */ private void assignPartition(TopicPartition partition, TreeSet sortedCurrentSubscriptions, Map> currentAssignment, - Map> consumer2AllPotentialPartitions, + Map> consumer2AllPotentialTopics, Map currentPartitionConsumer) { for (String consumer: sortedCurrentSubscriptions) { - if (consumer2AllPotentialPartitions.get(consumer).contains(partition)) { + if (consumer2AllPotentialTopics.get(consumer).contains(partition.topic())) { sortedCurrentSubscriptions.remove(consumer); currentAssignment.get(consumer).add(partition); currentPartitionConsumer.put(partition, consumer); @@ -637,19 +725,23 @@ private void assignPartition(TopicPartition partition, } } - private boolean canParticipateInReassignment(TopicPartition partition, - Map> partition2AllPotentialConsumers) { - // if a partition has two or more potential consumers it is subject to reassignment. - return partition2AllPotentialConsumers.get(partition).size() >= 2; + private boolean canParticipateInReassignment(String topic, + Map> topic2AllPotentialConsumers) { + // if a topic has two or more potential consumers it is subject to reassignment. + return topic2AllPotentialConsumers.get(topic).size() >= 2; } private boolean canParticipateInReassignment(String consumer, Map> currentAssignment, - Map> consumer2AllPotentialPartitions, - Map> partition2AllPotentialConsumers) { + Map> consumer2AllPotentialTopics, + Map> topic2AllPotentialConsumers, + Map partitionsPerTopic, + int totalPartitionCount) { List currentPartitions = currentAssignment.get(consumer); int currentAssignmentSize = currentPartitions.size(); - int maxAssignmentSize = consumer2AllPotentialPartitions.get(consumer).size(); + List allSubscribedTopics = consumer2AllPotentialTopics.get(consumer); + int maxAssignmentSize = getMaxAssignmentSize(totalPartitionCount, allSubscribedTopics, partitionsPerTopic); + if (currentAssignmentSize > maxAssignmentSize) log.error("The consumer {} is assigned more partitions than the maximum possible.", consumer); @@ -660,7 +752,7 @@ private boolean canParticipateInReassignment(String consumer, 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 (canParticipateInReassignment(partition, partition2AllPotentialConsumers)) + if (canParticipateInReassignment(partition.topic(), topic2AllPotentialConsumers)) return true; return false; @@ -674,36 +766,40 @@ private void balance(Map> currentAssignment, List sortedPartitions, List unassignedPartitions, TreeSet sortedCurrentSubscriptions, - Map> consumer2AllPotentialPartitions, - Map> partition2AllPotentialConsumers, + Map> consumer2AllPotentialTopics, + Map> topic2AllPotentialConsumers, Map currentPartitionConsumer, - boolean revocationRequired) { + boolean revocationRequired, + Map partitionsPerTopic, + int totalPartitionCount) { boolean initializing = currentAssignment.get(sortedCurrentSubscriptions.last()).isEmpty(); - boolean reassignmentPerformed = false; // assign all unassigned partitions for (TopicPartition partition: unassignedPartitions) { - // skip if there is no potential consumer for the partition - if (partition2AllPotentialConsumers.get(partition).isEmpty()) + // skip if there is no potential consumer for the topic + if (topic2AllPotentialConsumers.get(partition.topic()).isEmpty()) continue; assignPartition(partition, sortedCurrentSubscriptions, currentAssignment, - consumer2AllPotentialPartitions, currentPartitionConsumer); + consumer2AllPotentialTopics, currentPartitionConsumer); } // narrow down the reassignment scope to only those partitions that can actually be reassigned Set fixedPartitions = new HashSet<>(); - for (TopicPartition partition: partition2AllPotentialConsumers.keySet()) - if (!canParticipateInReassignment(partition, partition2AllPotentialConsumers)) - fixedPartitions.add(partition); + for (String topic: topic2AllPotentialConsumers.keySet()) + if (!canParticipateInReassignment(topic, topic2AllPotentialConsumers)) { + for (int i = 0; i < partitionsPerTopic.get(topic); i++) { + fixedPartitions.add(new TopicPartition(topic, i)); + } + } sortedPartitions.removeAll(fixedPartitions); unassignedPartitions.removeAll(fixedPartitions); // narrow down the reassignment scope to only those consumers that are subject to reassignment Map> fixedAssignments = new HashMap<>(); - for (String consumer: consumer2AllPotentialPartitions.keySet()) + for (String consumer: consumer2AllPotentialTopics.keySet()) if (!canParticipateInReassignment(consumer, currentAssignment, - consumer2AllPotentialPartitions, partition2AllPotentialConsumers)) { + consumer2AllPotentialTopics, topic2AllPotentialConsumers, partitionsPerTopic, totalPartitionCount)) { sortedCurrentSubscriptions.remove(consumer); fixedAssignments.put(consumer, currentAssignment.remove(consumer)); } @@ -715,11 +811,11 @@ private void balance(Map> currentAssignment, // 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, currentAssignment, prevAssignment, sortedCurrentSubscriptions, - consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer); + consumer2AllPotentialTopics, topic2AllPotentialConsumers, currentPartitionConsumer, partitionsPerTopic, totalPartitionCount); } - reassignmentPerformed = performReassignments(sortedPartitions, currentAssignment, prevAssignment, sortedCurrentSubscriptions, - consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer); + boolean reassignmentPerformed = performReassignments(sortedPartitions, currentAssignment, prevAssignment, sortedCurrentSubscriptions, + consumer2AllPotentialTopics, topic2AllPotentialConsumers, currentPartitionConsumer, partitionsPerTopic, totalPartitionCount); // 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 @@ -743,9 +839,11 @@ private boolean performReassignments(List reassignablePartitions Map> currentAssignment, Map prevAssignment, TreeSet sortedCurrentSubscriptions, - Map> consumer2AllPotentialPartitions, - Map> partition2AllPotentialConsumers, - Map currentPartitionConsumer) { + Map> consumer2AllPotentialTopics, + Map> topic2AllPotentialConsumers, + Map currentPartitionConsumer, + Map partitionsPerTopic, + int totalPartitionCount) { boolean reassignmentPerformed = false; boolean modified; @@ -755,11 +853,12 @@ private boolean performReassignments(List reassignablePartitions // 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 partitionIterator = reassignablePartitions.iterator(); - while (partitionIterator.hasNext() && !isBalanced(currentAssignment, sortedCurrentSubscriptions, consumer2AllPotentialPartitions)) { + while (partitionIterator.hasNext() && !isBalanced(currentAssignment, sortedCurrentSubscriptions, + consumer2AllPotentialTopics, partitionsPerTopic, totalPartitionCount)) { TopicPartition partition = partitionIterator.next(); // the partition must have at least two consumers - if (partition2AllPotentialConsumers.get(partition).size() <= 1) + 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 @@ -776,9 +875,9 @@ private boolean performReassignments(List reassignablePartitions } // check if a better-suited consumer exist for the partition; if so, reassign it - for (String otherConsumer: partition2AllPotentialConsumers.get(partition)) { + for (String otherConsumer: topic2AllPotentialConsumers.get(partition.topic())) { if (currentAssignment.get(consumer).size() > currentAssignment.get(otherConsumer).size() + 1) { - reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, consumer2AllPotentialPartitions); + reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, consumer2AllPotentialTopics); reassignmentPerformed = true; modified = true; break; @@ -794,11 +893,11 @@ private void reassignPartition(TopicPartition partition, Map> currentAssignment, TreeSet sortedCurrentSubscriptions, Map currentPartitionConsumer, - Map> consumer2AllPotentialPartitions) { + Map> consumer2AllPotentialTopics) { // find the new consumer String newConsumer = null; for (String anotherConsumer: sortedCurrentSubscriptions) { - if (consumer2AllPotentialPartitions.get(anotherConsumer).contains(partition)) { + if (consumer2AllPotentialTopics.get(anotherConsumer).contains(partition.topic())) { newConsumer = anotherConsumer; break; } @@ -855,17 +954,35 @@ private Map> deepCopy(Map, Serializable { + private static final long serialVersionUID = 1L; + private Map> map; + + TopicComparator(Map> 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 PartitionComparator implements Comparator, Serializable { private static final long serialVersionUID = 1L; - private Map> map; + private Map> map; - PartitionComparator(Map> map) { + PartitionComparator(Map> map) { this.map = map; } @Override public int compare(TopicPartition o1, TopicPartition o2) { - int ret = map.get(o1).size() - map.get(o2).size(); + int ret = map.get(o1.topic()).size() - map.get(o2.topic()).size(); if (ret == 0) { ret = o1.topic().compareTo(o2.topic()); if (ret == 0) diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java index edc522bb082ec..684a421be1807 100644 --- a/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java @@ -37,6 +37,8 @@ import org.apache.kafka.common.protocol.types.Struct; import org.apache.kafka.common.utils.CollectionUtils; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; public class StickyAssignorTest extends AbstractStickyAssignorTest { @@ -51,132 +53,147 @@ public Subscription buildSubscription(List topics, List serializeTopicPartitionAssignment(new MemberData(partitions, Optional.of(DEFAULT_GENERATION)))); } - @Test - public void testAssignmentWithMultipleGenerations1() { - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - String consumer3 = "consumer3"; + @ParameterizedTest(name = "testAssignmentWithMultipleGenerations1 with isAllSubscriptionsEqual: {0}") + @ValueSource(booleans = {true, false}) + public void testAssignmentWithMultipleGenerations1(boolean isAllSubscriptionsEqual) { + List allTopics = topics(topic, topic2); + List consumer2SubscribedTopics = isAllSubscriptionsEqual ? allTopics : topics(topic); Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic, 6); - subscriptions.put(consumer1, new Subscription(topics(topic))); - subscriptions.put(consumer2, new Subscription(topics(topic))); - subscriptions.put(consumer3, new Subscription(topics(topic))); + partitionsPerTopic.put(topic2, 6); + subscriptions.put(consumer1, new Subscription(allTopics)); + subscriptions.put(consumer2, new Subscription(consumer2SubscribedTopics)); + subscriptions.put(consumer3, new Subscription(allTopics)); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); List r1partitions1 = assignment.get(consumer1); List r1partitions2 = assignment.get(consumer2); List r1partitions3 = assignment.get(consumer3); - assertTrue(r1partitions1.size() == 2 && r1partitions2.size() == 2 && r1partitions3.size() == 2); + assertTrue(r1partitions1.size() == 4 && r1partitions2.size() == 4 && r1partitions3.size() == 4); verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic); assertTrue(isFullyBalanced(assignment)); - subscriptions.put(consumer1, buildSubscription(topics(topic), r1partitions1)); - subscriptions.put(consumer2, buildSubscription(topics(topic), r1partitions2)); + subscriptions.put(consumer1, buildSubscription(allTopics, r1partitions1)); + subscriptions.put(consumer2, buildSubscription(consumer2SubscribedTopics, r1partitions2)); subscriptions.remove(consumer3); assignment = assignor.assign(partitionsPerTopic, subscriptions); List r2partitions1 = assignment.get(consumer1); List r2partitions2 = assignment.get(consumer2); - assertTrue(r2partitions1.size() == 3 && r2partitions2.size() == 3); - assertTrue(r2partitions1.containsAll(r1partitions1)); + assertTrue(r2partitions1.size() == 6 && r2partitions2.size() == 6); + if (isAllSubscriptionsEqual) { + // only true in all subscription equal case + assertTrue(r2partitions1.containsAll(r1partitions1)); + } assertTrue(r2partitions2.containsAll(r1partitions2)); verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic); assertTrue(isFullyBalanced(assignment)); assertFalse(Collections.disjoint(r2partitions2, r1partitions3)); subscriptions.remove(consumer1); - subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic), r2partitions2, 2)); - subscriptions.put(consumer3, buildSubscriptionWithGeneration(topics(topic), r1partitions3, 1)); + subscriptions.put(consumer2, buildSubscriptionWithGeneration(consumer2SubscribedTopics, r2partitions2, 2)); + subscriptions.put(consumer3, buildSubscriptionWithGeneration(allTopics, r1partitions3, 1)); assignment = assignor.assign(partitionsPerTopic, subscriptions); List r3partitions2 = assignment.get(consumer2); List r3partitions3 = assignment.get(consumer3); - assertTrue(r3partitions2.size() == 3 && r3partitions3.size() == 3); + assertTrue(r3partitions2.size() == 6 && r3partitions3.size() == 6); assertTrue(Collections.disjoint(r3partitions2, r3partitions3)); verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic); assertTrue(isFullyBalanced(assignment)); } - @Test - public void testAssignmentWithMultipleGenerations2() { - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - String consumer3 = "consumer3"; + @ParameterizedTest(name = "testAssignmentWithMultipleGenerations2 with isAllSubscriptionsEqual: {0}") + @ValueSource(booleans = {true, false}) + public void testAssignmentWithMultipleGenerations2(boolean isAllSubscriptionsEqual) { + List allTopics = topics(topic, topic2, topic3); + List consumer1SubscribedTopics = isAllSubscriptionsEqual ? allTopics : topics(topic); + List consumer3SubscribedTopics = isAllSubscriptionsEqual ? allTopics : topics(topic, topic2); Map partitionsPerTopic = new HashMap<>(); - partitionsPerTopic.put(topic, 6); - subscriptions.put(consumer1, new Subscription(topics(topic))); - subscriptions.put(consumer2, new Subscription(topics(topic))); - subscriptions.put(consumer3, new Subscription(topics(topic))); + partitionsPerTopic.put(topic, 4); + partitionsPerTopic.put(topic2, 4); + partitionsPerTopic.put(topic3, 4); + subscriptions.put(consumer1, new Subscription(consumer1SubscribedTopics)); + subscriptions.put(consumer2, new Subscription(allTopics)); + subscriptions.put(consumer3, new Subscription(consumer3SubscribedTopics)); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); List r1partitions1 = assignment.get(consumer1); List r1partitions2 = assignment.get(consumer2); List r1partitions3 = assignment.get(consumer3); - assertTrue(r1partitions1.size() == 2 && r1partitions2.size() == 2 && r1partitions3.size() == 2); + assertTrue(r1partitions1.size() == 4 && r1partitions2.size() == 4 && r1partitions3.size() == 4); verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic); assertTrue(isFullyBalanced(assignment)); subscriptions.remove(consumer1); - subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic), r1partitions2, 1)); + subscriptions.put(consumer2, buildSubscriptionWithGeneration(allTopics, r1partitions2, 1)); subscriptions.remove(consumer3); assignment = assignor.assign(partitionsPerTopic, subscriptions); List r2partitions2 = assignment.get(consumer2); - assertEquals(6, r2partitions2.size()); + assertEquals(12, r2partitions2.size()); assertTrue(r2partitions2.containsAll(r1partitions2)); verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic); assertTrue(isFullyBalanced(assignment)); - subscriptions.put(consumer1, buildSubscriptionWithGeneration(topics(topic), r1partitions1, 1)); - subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic), r2partitions2, 2)); - subscriptions.put(consumer3, buildSubscriptionWithGeneration(topics(topic), r1partitions3, 1)); + subscriptions.put(consumer1, buildSubscriptionWithGeneration(consumer1SubscribedTopics, r1partitions1, 1)); + subscriptions.put(consumer2, buildSubscriptionWithGeneration(allTopics, r2partitions2, 2)); + subscriptions.put(consumer3, buildSubscriptionWithGeneration(consumer3SubscribedTopics, r1partitions3, 1)); assignment = assignor.assign(partitionsPerTopic, subscriptions); List r3partitions1 = assignment.get(consumer1); List r3partitions2 = assignment.get(consumer2); List r3partitions3 = assignment.get(consumer3); - assertTrue(r3partitions1.size() == 2 && r3partitions2.size() == 2 && r3partitions3.size() == 2); - assertEquals(r1partitions1, r3partitions1); - assertEquals(r1partitions2, r3partitions2); - assertEquals(r1partitions3, r3partitions3); + assertTrue(r3partitions1.size() == 4 && r3partitions2.size() == 4 && r3partitions3.size() == 4); verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic); assertTrue(isFullyBalanced(assignment)); } - @Test - public void testAssignmentWithConflictingPreviousGenerations() { - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - String consumer3 = "consumer3"; - + @ParameterizedTest(name = "testAssignmentWithConflictingPreviousGenerations with isAllSubscriptionsEqual: {0}") + @ValueSource(booleans = {true, false}) + public void testAssignmentWithConflictingPreviousGenerations(boolean isAllSubscriptionsEqual) { Map partitionsPerTopic = new HashMap<>(); - partitionsPerTopic.put(topic, 6); - subscriptions.put(consumer1, new Subscription(topics(topic))); - subscriptions.put(consumer2, new Subscription(topics(topic))); - subscriptions.put(consumer3, new Subscription(topics(topic))); + partitionsPerTopic.put(topic, 4); + partitionsPerTopic.put(topic2, 4); + partitionsPerTopic.put(topic3, 4); + + List allTopics = topics(topic, topic2, topic3); + List consumer1SubscribedTopics = isAllSubscriptionsEqual ? allTopics : topics(topic); + List consumer2SubscribedTopics = isAllSubscriptionsEqual ? allTopics : topics(topic, topic2); + + subscriptions.put(consumer1, new Subscription(consumer1SubscribedTopics)); + subscriptions.put(consumer2, new Subscription(consumer2SubscribedTopics)); + subscriptions.put(consumer3, new Subscription(allTopics)); TopicPartition tp0 = new TopicPartition(topic, 0); TopicPartition tp1 = new TopicPartition(topic, 1); TopicPartition tp2 = new TopicPartition(topic, 2); TopicPartition tp3 = new TopicPartition(topic, 3); - TopicPartition tp4 = new TopicPartition(topic, 4); - TopicPartition tp5 = new TopicPartition(topic, 5); - - List c1partitions0 = partitions(tp0, tp1, tp4); - List c2partitions0 = partitions(tp0, tp1, tp2); - List c3partitions0 = partitions(tp3, tp4, tp5); - subscriptions.put(consumer1, buildSubscriptionWithGeneration(topics(topic), c1partitions0, 1)); - subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic), c2partitions0, 2)); - subscriptions.put(consumer3, buildSubscriptionWithGeneration(topics(topic), c3partitions0, 2)); + TopicPartition t2p0 = new TopicPartition(topic2, 0); + TopicPartition t2p1 = new TopicPartition(topic2, 1); + TopicPartition t2p2 = new TopicPartition(topic2, 2); + TopicPartition t2p3 = new TopicPartition(topic2, 3); + TopicPartition t3p0 = new TopicPartition(topic3, 0); + TopicPartition t3p1 = new TopicPartition(topic3, 1); + TopicPartition t3p2 = new TopicPartition(topic3, 2); + TopicPartition t3p3 = new TopicPartition(topic3, 3); + + List c1partitions0 = isAllSubscriptionsEqual ? partitions(tp0, tp1, tp2, t2p2, t2p3, t3p0) : + partitions(tp0, tp1, tp2, tp3); + List c2partitions0 = partitions(tp0, tp1, t2p0, t2p1, t2p2, t2p3); + List c3partitions0 = partitions(tp2, tp3, t3p0, t3p1, t3p2, t3p3); + subscriptions.put(consumer1, buildSubscriptionWithGeneration(consumer1SubscribedTopics, c1partitions0, 1)); + subscriptions.put(consumer2, buildSubscriptionWithGeneration(consumer2SubscribedTopics, c2partitions0, 2)); + subscriptions.put(consumer3, buildSubscriptionWithGeneration(allTopics, c3partitions0, 2)); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); List c3partitions = assignment.get(consumer3); - assertTrue(c1partitions.size() == 2 && c2partitions.size() == 2 && c3partitions.size() == 2); + assertTrue(c1partitions.size() == 4 && c2partitions.size() == 4 && c3partitions.size() == 4); assertTrue(c2partitions0.containsAll(c2partitions)); assertTrue(c3partitions0.containsAll(c3partitions)); verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic); @@ -185,10 +202,6 @@ public void testAssignmentWithConflictingPreviousGenerations() { @Test public void testSchemaBackwardCompatibility() { - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - String consumer3 = "consumer3"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic, 3); subscriptions.put(consumer1, new Subscription(topics(topic))); diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java index c8f6c14c5e82a..a650cbbf39a67 100644 --- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java @@ -43,8 +43,15 @@ public abstract class AbstractStickyAssignorTest { protected AbstractStickyAssignor assignor; protected String consumerId = "consumer"; + protected String consumer1 = "consumer1"; + protected String consumer2 = "consumer2"; + protected String consumer3 = "consumer3"; + protected String consumer4 = "consumer4"; protected Map subscriptions; protected String topic = "topic"; + protected String topic1 = "topic1"; + protected String topic2 = "topic2"; + protected String topic3 = "topic3"; protected abstract AbstractStickyAssignor createAssignor(); @@ -124,9 +131,6 @@ public void testOnlyAssignsPartitionsFromSubscribedTopics() { @Test public void testOneConsumerMultipleTopics() { - String topic1 = "topic1"; - String topic2 = "topic2"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic1, 1); partitionsPerTopic.put(topic2, 2); @@ -141,9 +145,6 @@ public void testOneConsumerMultipleTopics() { @Test public void testTwoConsumersOneTopicOnePartition() { - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic, 1); @@ -158,9 +159,6 @@ public void testTwoConsumersOneTopicOnePartition() { @Test public void testTwoConsumersOneTopicTwoPartitions() { - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic, 2); @@ -177,12 +175,6 @@ public void testTwoConsumersOneTopicTwoPartitions() { @Test public void testMultipleConsumersMixedTopicSubscriptions() { - String topic1 = "topic1"; - String topic2 = "topic2"; - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - String consumer3 = "consumer3"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic1, 3); partitionsPerTopic.put(topic2, 2); @@ -202,11 +194,6 @@ public void testMultipleConsumersMixedTopicSubscriptions() { @Test public void testTwoConsumersTwoTopicsSixPartitions() { - String topic1 = "topic1"; - String topic2 = "topic2"; - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic1, 3); partitionsPerTopic.put(topic2, 3); @@ -227,11 +214,6 @@ public void testTwoConsumersTwoTopicsSixPartitions() { */ @Test public void testConsumerOwningMinQuotaExpectedMaxQuota() { - String topic1 = "topic1"; - String topic2 = "topic2"; - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic1, 2); partitionsPerTopic.put(topic2, 3); @@ -256,12 +238,6 @@ public void testConsumerOwningMinQuotaExpectedMaxQuota() { */ @Test public void testMaxQuotaConsumerMoreThanNumExpectedMaxCapacityMembers() { - String topic1 = "topic1"; - String topic2 = "topic2"; - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - String consumer3 = "consumer3"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic1, 2); partitionsPerTopic.put(topic2, 2); @@ -289,12 +265,6 @@ public void testMaxQuotaConsumerMoreThanNumExpectedMaxCapacityMembers() { */ @Test public void testAllConsumerAreUnderMinQuota() { - String topic1 = "topic1"; - String topic2 = "topic2"; - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - String consumer3 = "consumer3"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic1, 2); partitionsPerTopic.put(topic2, 3); @@ -319,8 +289,6 @@ public void testAllConsumerAreUnderMinQuota() { @Test public void testAddRemoveConsumerOneTopic() { - String consumer1 = "consumer1"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic, 3); subscriptions.put(consumer1, new Subscription(topics(topic))); @@ -331,7 +299,6 @@ public void testAddRemoveConsumerOneTopic() { verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic); assertTrue(isFullyBalanced(assignment)); - String consumer2 = "consumer2"; subscriptions.put(consumer1, buildSubscription(topics(topic), assignment.get(consumer1))); subscriptions.put(consumer2, buildSubscription(topics(topic), Collections.emptyList())); assignment = assignor.assign(partitionsPerTopic, subscriptions); @@ -353,12 +320,6 @@ public void testAddRemoveConsumerOneTopic() { @Test public void testAddRemoveTwoConsumersTwoTopics() { - String topic1 = "topic1"; - String topic2 = "topic2"; - String consumer1 = "consumer1"; - String consumer2 = "consumer2"; - String consumer3 = "consumer3"; - String consumer4 = "consumer4"; List allTopics = topics(topic1, topic2); Map partitionsPerTopic = new HashMap<>(); @@ -441,9 +402,6 @@ public void testPoorRoundRobinAssignmentScenario() { @Test public void testAddRemoveTopicTwoConsumers() { - String consumer1 = "consumer"; - String consumer2 = "consumer2"; - Map partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic, 3); subscriptions.put(consumer1, new Subscription(topics(topic))); @@ -459,7 +417,6 @@ public void testAddRemoveTopicTwoConsumers() { assertTrue((consumer1Assignment1.size() == 1 && consumer2Assignment1.size() == 2) || (consumer1Assignment1.size() == 2 && consumer2Assignment1.size() == 1)); - String topic2 = "topic2"; partitionsPerTopic.put(topic2, 3); subscriptions.put(consumer1, buildSubscription(topics(topic, topic2), assignment.get(consumer1))); subscriptions.put(consumer2, buildSubscription(topics(topic, topic2), assignment.get(consumer2))); @@ -598,6 +555,43 @@ public void testLargeAssignmentAndGroupWithUniformSubscription() { assignor.assign(partitionsPerTopic, subscriptions); } + @Timeout(40) + @Test + public void testLargeAssignmentAndGroupWithNonEqualSubscription() { + // 1 million partitions! + int topicCount = 500; + int partitionCount = 2_000; + int consumerCount = 2_000; + + List topics = new ArrayList<>(); + Map partitionsPerTopic = new HashMap<>(); + for (int i = 0; i < topicCount; i++) { + String topicName = getTopicName(i, topicCount); + topics.add(topicName); + partitionsPerTopic.put(topicName, partitionCount); + } + for (int i = 0; i < consumerCount; i++) { + if (i == consumerCount - 1) { + subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics.subList(0, 1))); + } else { + subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics)); + } + } + + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); + + for (int i = 1; i < consumerCount; i++) { + String consumer = getConsumerName(i, consumerCount); + if (i == consumerCount - 1) { + subscriptions.put(consumer, buildSubscription(topics.subList(0, 1), assignment.get(consumer))); + } else { + subscriptions.put(consumer, buildSubscription(topics, assignment.get(consumer))); + } + } + + assignor.assign(partitionsPerTopic, subscriptions); + } + @Test public void testLargeAssignmentWithMultipleConsumersLeavingAndRandomSubscription() { Random rand = new Random(); @@ -658,19 +652,23 @@ public void testNewSubscription() { @Test public void testMoveExistingAssignments() { + String topic4 = "topic4"; + String topic5 = "topic5"; + String topic6 = "topic6"; + Map partitionsPerTopic = new HashMap<>(); for (int i = 1; i <= 6; i++) - partitionsPerTopic.put(String.format("topic%02d", i), 1); - - subscriptions.put("consumer01", - buildSubscription(topics("topic01", "topic02"), - partitions(tp("topic01", 0)))); - subscriptions.put("consumer02", - buildSubscription(topics("topic01", "topic02", "topic03", "topic04"), - partitions(tp("topic02", 0), tp("topic03", 0)))); - subscriptions.put("consumer03", - buildSubscription(topics("topic02", "topic03", "topic04", "topic05", "topic06"), - partitions(tp("topic04", 0), tp("topic05", 0), tp("topic06", 0)))); + partitionsPerTopic.put(String.format("topic%d", i), 1); + + subscriptions.put(consumer1, + buildSubscription(topics(topic1, topic2), + partitions(tp(topic1, 0)))); + subscriptions.put(consumer2, + buildSubscription(topics(topic1, topic2, topic3, topic4), + partitions(tp(topic2, 0), tp(topic3, 0)))); + subscriptions.put(consumer3, + buildSubscription(topics(topic2, topic3, topic4, topic5, topic6), + partitions(tp(topic4, 0), tp(topic5, 0), tp(topic6, 0)))); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic); @@ -679,16 +677,12 @@ public void testMoveExistingAssignments() { @Test public void testStickiness() { Map partitionsPerTopic = new HashMap<>(); - partitionsPerTopic.put("topic01", 3); - String consumer1 = "consumer01"; - String consumer2 = "consumer02"; - String consumer3 = "consumer03"; - String consumer4 = "consumer04"; + partitionsPerTopic.put(topic1, 3); - subscriptions.put(consumer1, new Subscription(topics("topic01"))); - subscriptions.put(consumer2, new Subscription(topics("topic01"))); - subscriptions.put(consumer3, new Subscription(topics("topic01"))); - subscriptions.put(consumer4, new Subscription(topics("topic01"))); + subscriptions.put(consumer1, new Subscription(topics(topic1))); + subscriptions.put(consumer2, new Subscription(topics(topic1))); + subscriptions.put(consumer3, new Subscription(topics(topic1))); + subscriptions.put(consumer4, new Subscription(topics(topic1))); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic); @@ -707,11 +701,11 @@ public void testStickiness() { // removing the potential group leader subscriptions.remove(consumer1); subscriptions.put(consumer2, - buildSubscription(topics("topic01"), assignment.get(consumer2))); + buildSubscription(topics(topic1), assignment.get(consumer2))); subscriptions.put(consumer3, - buildSubscription(topics("topic01"), assignment.get(consumer3))); + buildSubscription(topics(topic1), assignment.get(consumer3))); subscriptions.put(consumer4, - buildSubscription(topics("topic01"), assignment.get(consumer4))); + buildSubscription(topics(topic1), assignment.get(consumer4))); assignment = assignor.assign(partitionsPerTopic, subscriptions); @@ -730,9 +724,9 @@ public void testStickiness() { @Test public void testAssignmentUpdatedForDeletedTopic() { Map partitionsPerTopic = new HashMap<>(); - partitionsPerTopic.put("topic01", 1); - partitionsPerTopic.put("topic03", 100); - subscriptions = Collections.singletonMap(consumerId, new Subscription(topics("topic01", "topic02", "topic03"))); + partitionsPerTopic.put(topic1, 1); + partitionsPerTopic.put(topic3, 100); + subscriptions = Collections.singletonMap(consumerId, new Subscription(topics(topic1, topic2, topic3))); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); assertEquals(assignment.values().stream().mapToInt(List::size).sum(), 1 + 100); @@ -764,12 +758,12 @@ public void testReassignmentWithRandomSubscriptionsAndChanges() { int numTopics = minNumTopics + new Random().nextInt(maxNumTopics - minNumTopics); ArrayList topics = new ArrayList<>(); - for (int i = 0; i < numTopics; ++i) - topics.add(getTopicName(i, maxNumTopics)); Map partitionsPerTopic = new HashMap<>(); - for (int i = 0; i < numTopics; ++i) + for (int i = 0; i < numTopics; ++i) { + topics.add(getTopicName(i, maxNumTopics)); partitionsPerTopic.put(getTopicName(i, maxNumTopics), i + 1); + } int numConsumers = minNumConsumers + new Random().nextInt(maxNumConsumers - minNumConsumers);