You are viewing a plain text version of this content. The canonical link for it is here.
Posted to jira@kafka.apache.org by GitBox <gi...@apache.org> on 2021/05/03 22:59:13 UTC
[GitHub] [kafka] guozhangwang commented on a change in pull request #10552: KAFKA-12675: improve the sticky general assignor scalability and performance
guozhangwang commented on a change in pull request #10552:
URL: https://github.com/apache/kafka/pull/10552#discussion_r625409879
##########
File path: clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
##########
@@ -307,32 +306,35 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
*/
private Map<String, List<TopicPartition>> generalAssign(Map<String, Integer> partitionsPerTopic,
Map<String, Subscription> subscriptions) {
+ if (log.isDebugEnabled()) {
+ log.debug(String.format("performing general assign. partitionsPerTopic: %s, subscriptions: %s",
+ partitionsPerTopic, subscriptions));
+ }
+
Map<String, List<TopicPartition>> currentAssignment = new HashMap<>();
Map<TopicPartition, ConsumerGenerationPair> prevAssignment = new HashMap<>();
partitionMovements = new PartitionMovements();
prepopulateCurrentAssignments(subscriptions, currentAssignment, prevAssignment);
- // a mapping of all topic partitions to all consumers that can be assigned to them
- final Map<TopicPartition, List<String>> partition2AllPotentialConsumers = new HashMap<>();
- // a mapping of all consumers to all potential topic partitions that can be assigned to them
- final Map<String, List<TopicPartition>> consumer2AllPotentialPartitions = new HashMap<>();
+ // a mapping of all topics to all consumers that can be assigned to them
+ final Map<String, List<String>> topic2AllPotentialConsumers = new HashMap<>(partitionsPerTopic.keySet().size());
+ // a mapping of all consumers to all potential topics that can be assigned to them
+ final Map<String, List<String>> consumer2AllPotentialTopics = new HashMap<>(subscriptions.keySet().size());
- // initialize partition2AllPotentialConsumers and consumer2AllPotentialPartitions in the following two for loops
+ // initialize topic2AllPotentialConsumers and consumer2AllPotentialTopics in the following two for loops
for (Entry<String, Integer> entry: partitionsPerTopic.entrySet()) {
for (int i = 0; i < entry.getValue(); ++i)
- partition2AllPotentialConsumers.put(new TopicPartition(entry.getKey(), i), new ArrayList<>());
+ topic2AllPotentialConsumers.put(entry.getKey(), new ArrayList<>());
Review comment:
Hmm is this right? Wouldn't we put the same empty list for the key N times?
##########
File path: clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
##########
@@ -387,58 +403,125 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
TreeSet<String> sortedCurrentSubscriptions = new TreeSet<>(new SubscriptionComparator(currentAssignment));
sortedCurrentSubscriptions.addAll(currentAssignment.keySet());
+ int totalPartitionCount = partitionsPerTopic.values().stream().reduce(0, Integer::sum);
+
balance(currentAssignment, prevAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions,
- consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer, revocationRequired);
+ consumer2AllPotentialTopics, topic2AllPotentialConsumers, currentPartitionConsumer, revocationRequired,
+ partitionsPerTopic, totalPartitionCount);
+
+ 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 sortedToBeRemovedPartitions. We use two pointers technique here:
+ *
+ * We loop the sortedPartition, and compare the ith element in sorted toBeRemovedPartitions(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 sortedToBeRemovedPartitions
+ *
+ * @param sortedPartitions: sorted all partitions
+ * @param sortedToBeRemovedPartitions: sorted partitions, all are included in the sortedPartitions
+ * @return the partitions don't assign to any current consumers
+ */
+ private List<TopicPartition> getUnassignedPartitions(List<TopicPartition> sortedPartitions,
+ List<TopicPartition> sortedToBeRemovedPartitions) {
+ List<TopicPartition> unassignedPartitions = new ArrayList<>();
+
+ int index = 0;
+ boolean shouldAddDirectly = false;
+ int sizeToBeRemovedPartitions = sortedToBeRemovedPartitions.size();
+ TopicPartition nextPartition = sortedToBeRemovedPartitions.get(index);
+ for (TopicPartition topicPartition : sortedPartitions) {
+ if (shouldAddDirectly || !nextPartition.equals(topicPartition)) {
+ unassignedPartitions.add(topicPartition);
+ } else {
+ // equal case, don't add to unassignedPartitions, just get next partition
+ if (index < sizeToBeRemovedPartitions - 1) {
+ nextPartition = sortedToBeRemovedPartitions.get(++index);
+ } else {
+ // add the remaining directly since there is no more toBeRemovedPartitions
+ 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<TopicPartition, ConsumerGenerationPair> prevAssignment,
+ List<TopicPartition> partitions,
+ String consumer,
+ int generation) {
+ for (TopicPartition partition: partitions) {
+ ConsumerGenerationPair consumerGeneration = prevAssignment.get(partition);
+ if (consumerGeneration != null) {
+ // only keep the latest previous assignment
+ if (generation > consumerGeneration.generation)
+ prevAssignment.put(partition, new ConsumerGenerationPair(consumer, generation));
+ } else {
+ prevAssignment.put(partition, new ConsumerGenerationPair(consumer, generation));
+ }
+ }
+ }
+
+ /**
+ * filling in the currentAssignment and prevAssignment from the subscriptions.
+ *
+ * @param subscriptions: Map from the member id to their respective topic subscription
+ * @param currentAssignment: The assignment contains the assignments with the largest generation
+ * @param prevAssignment: The assignment contains the assignment with the 2nd largest generation
+ */
private void prepopulateCurrentAssignments(Map<String, Subscription> subscriptions,
Map<String, List<TopicPartition>> currentAssignment,
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 exists only if user data is for different generations
- // for each partition we create a sorted map of its consumers by generation
- Map<TopicPartition, TreeMap<Integer, String>> sortedPartitionConsumersByGeneration = new HashMap<>();
+ Set<String> membersOfCurrentHighestGeneration = new HashSet<>();
+ int maxGeneration = DEFAULT_GENERATION;
+
for (Map.Entry<String, Subscription> subscriptionEntry: subscriptions.entrySet()) {
String consumer = subscriptionEntry.getKey();
MemberData memberData = memberData(subscriptionEntry.getValue());
- for (TopicPartition partition: memberData.partitions) {
- if (sortedPartitionConsumersByGeneration.containsKey(partition)) {
- Map<Integer, String> 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<Integer, String> sortedConsumers = new TreeMap<>();
- sortedConsumers.put(memberData.generation.orElse(DEFAULT_GENERATION), consumer);
- sortedPartitionConsumersByGeneration.put(partition, sortedConsumers);
- }
- }
- }
+ List<TopicPartition> ownedPartitions = new ArrayList<>();
+ currentAssignment.put(consumer, ownedPartitions);
- // 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<TopicPartition, TreeMap<Integer, String>> partitionConsumersEntry: sortedPartitionConsumersByGeneration.entrySet()) {
- TopicPartition partition = partitionConsumersEntry.getKey();
- TreeMap<Integer, String> consumers = partitionConsumersEntry.getValue();
- Iterator<Integer> it = consumers.descendingKeySet().iterator();
+ // Only consider this consumer's owned partitions as valid if it is a member of the current highest
+ // generation, or it's generation is not present but we have not seen any known generation so far
+ if (memberData.generation.isPresent() && memberData.generation.get() >= maxGeneration
+ || !memberData.generation.isPresent() && maxGeneration == DEFAULT_GENERATION) {
- // 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);
+ // If the current member's generation is higher, all the previously owned partitions are invalid
+ if (memberData.generation.isPresent() && memberData.generation.get() > maxGeneration) {
+ for (String member: membersOfCurrentHighestGeneration) {
+ List<TopicPartition> oldGenerationPartitions = currentAssignment.get(member);
+ updatePrevAssignment(prevAssignment, oldGenerationPartitions, member, maxGeneration);
+ oldGenerationPartitions.clear();
+ }
+ membersOfCurrentHighestGeneration.clear();
+ maxGeneration = memberData.generation.get();
+ }
- // now update previous assignment if any
- if (it.hasNext()) {
- int generation = it.next();
- prevAssignment.put(partition, new ConsumerGenerationPair(consumers.get(generation), generation));
+ membersOfCurrentHighestGeneration.add(consumer);
+ ownedPartitions.addAll(memberData.partitions);
+ } else if (!memberData.generation.isPresent()) {
+ // current maxGeneration is larger than DEFAULT_GENERATION,
+ // put all partitions as DEFAULT_GENERATION into provAssignment
Review comment:
typo: prev.
##########
File path: clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
##########
@@ -362,23 +365,36 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
// otherwise (the consumer still exists)
for (Iterator<TopicPartition> partitionIter = entry.getValue().iterator(); partitionIter.hasNext();) {
TopicPartition partition = partitionIter.next();
- if (!partition2AllPotentialConsumers.containsKey(partition)) {
+ if (!topic2AllPotentialConsumers.containsKey(partition.topic())) {
Review comment:
nit: the following comment needs to be updated as well.
##########
File path: clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
##########
@@ -527,12 +615,23 @@ private int getBalanceScore(Map<String, List<TopicPartition>> assignment) {
* 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
+ * @param topic2AllPotentialConsumers a mapping of partitions to their potential consumers
+ * @param partitionsPerTopic The number of partitions for each subscribed topic
* @return an ascending sorted list of topic partitions based on how many consumers can potentially use them
*/
- private List<TopicPartition> sortPartitions(Map<TopicPartition, List<String>> partition2AllPotentialConsumers) {
- List<TopicPartition> sortedPartitions = new ArrayList<>(partition2AllPotentialConsumers.keySet());
- Collections.sort(sortedPartitions, new PartitionComparator(partition2AllPotentialConsumers));
+ private List<TopicPartition> sortPartitions(Map<String, List<String>> topic2AllPotentialConsumers,
+ Map<String, Integer> partitionsPerTopic) {
+ List<TopicPartition> sortedPartitions = new ArrayList<>();
+ List<String> allTopics = new ArrayList<>(topic2AllPotentialConsumers.keySet());
+ Collections.sort(allTopics, new TopicComparator(topic2AllPotentialConsumers));
+
+ // since allTopics are sorted, we can loop through allTopics to create the sortedPartitions
Review comment:
SGTM.
##########
File path: clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
##########
@@ -387,58 +398,121 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
TreeSet<String> sortedCurrentSubscriptions = new TreeSet<>(new SubscriptionComparator(currentAssignment));
sortedCurrentSubscriptions.addAll(currentAssignment.keySet());
+ int totalPartitionCount = partitionsPerTopic.values().stream().reduce(0, Integer::sum);
+
balance(currentAssignment, prevAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions,
- consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer, revocationRequired);
+ consumer2AllPotentialTopics, topic2AllPotentialConsumers, currentPartitionConsumer, revocationRequired,
+ partitionsPerTopic, totalPartitionCount);
+
return currentAssignment;
}
+ /**
+ * get the unassigned partition list by computing the difference set of the sortedPartitions(all partitions)
+ * and toBeRemovedPartitions. We use two pointers technique here:
+ *
+ * We loop the sortedPartition, and compare the ith element in sorted toBeRemovedPartitions(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 toBeRemovedPartitions
+ *
+ * @param sortedPartitions: sorted all partitions
+ * @param toBeRemovedPartitions: sorted partitions, all are included in the sortedPartitions
+ * @return the partitions don't assign to any current consumers
+ */
+ private List<TopicPartition> getUnassignedPartitions(List<TopicPartition> sortedPartitions,
Review comment:
Sounds great!
##########
File path: clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
##########
@@ -387,58 +403,125 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
TreeSet<String> sortedCurrentSubscriptions = new TreeSet<>(new SubscriptionComparator(currentAssignment));
sortedCurrentSubscriptions.addAll(currentAssignment.keySet());
+ int totalPartitionCount = partitionsPerTopic.values().stream().reduce(0, Integer::sum);
+
balance(currentAssignment, prevAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions,
- consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer, revocationRequired);
+ consumer2AllPotentialTopics, topic2AllPotentialConsumers, currentPartitionConsumer, revocationRequired,
+ partitionsPerTopic, totalPartitionCount);
+
+ 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 sortedToBeRemovedPartitions. We use two pointers technique here:
+ *
+ * We loop the sortedPartition, and compare the ith element in sorted toBeRemovedPartitions(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 sortedToBeRemovedPartitions
+ *
+ * @param sortedPartitions: sorted all partitions
+ * @param sortedToBeRemovedPartitions: sorted partitions, all are included in the sortedPartitions
+ * @return the partitions don't assign to any current consumers
+ */
+ private List<TopicPartition> getUnassignedPartitions(List<TopicPartition> sortedPartitions,
+ List<TopicPartition> sortedToBeRemovedPartitions) {
+ List<TopicPartition> unassignedPartitions = new ArrayList<>();
+
+ int index = 0;
+ boolean shouldAddDirectly = false;
+ int sizeToBeRemovedPartitions = sortedToBeRemovedPartitions.size();
+ TopicPartition nextPartition = sortedToBeRemovedPartitions.get(index);
+ for (TopicPartition topicPartition : sortedPartitions) {
+ if (shouldAddDirectly || !nextPartition.equals(topicPartition)) {
+ unassignedPartitions.add(topicPartition);
+ } else {
+ // equal case, don't add to unassignedPartitions, just get next partition
+ if (index < sizeToBeRemovedPartitions - 1) {
+ nextPartition = sortedToBeRemovedPartitions.get(++index);
+ } else {
+ // add the remaining directly since there is no more toBeRemovedPartitions
+ 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<TopicPartition, ConsumerGenerationPair> prevAssignment,
+ List<TopicPartition> partitions,
+ String consumer,
+ int generation) {
+ for (TopicPartition partition: partitions) {
+ ConsumerGenerationPair consumerGeneration = prevAssignment.get(partition);
+ if (consumerGeneration != null) {
+ // only keep the latest previous assignment
+ if (generation > consumerGeneration.generation)
+ prevAssignment.put(partition, new ConsumerGenerationPair(consumer, generation));
+ } else {
+ prevAssignment.put(partition, new ConsumerGenerationPair(consumer, generation));
+ }
+ }
+ }
+
+ /**
+ * filling in the currentAssignment and prevAssignment from the subscriptions.
+ *
+ * @param subscriptions: Map from the member id to their respective topic subscription
+ * @param currentAssignment: The assignment contains the assignments with the largest generation
+ * @param prevAssignment: The assignment contains the assignment with the 2nd largest generation
+ */
private void prepopulateCurrentAssignments(Map<String, Subscription> subscriptions,
Map<String, List<TopicPartition>> currentAssignment,
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 exists only if user data is for different generations
- // for each partition we create a sorted map of its consumers by generation
- Map<TopicPartition, TreeMap<Integer, String>> sortedPartitionConsumersByGeneration = new HashMap<>();
+ Set<String> membersOfCurrentHighestGeneration = new HashSet<>();
+ int maxGeneration = DEFAULT_GENERATION;
+
for (Map.Entry<String, Subscription> subscriptionEntry: subscriptions.entrySet()) {
String consumer = subscriptionEntry.getKey();
MemberData memberData = memberData(subscriptionEntry.getValue());
- for (TopicPartition partition: memberData.partitions) {
- if (sortedPartitionConsumersByGeneration.containsKey(partition)) {
- Map<Integer, String> 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<Integer, String> sortedConsumers = new TreeMap<>();
- sortedConsumers.put(memberData.generation.orElse(DEFAULT_GENERATION), consumer);
- sortedPartitionConsumersByGeneration.put(partition, sortedConsumers);
- }
- }
- }
+ List<TopicPartition> ownedPartitions = new ArrayList<>();
+ currentAssignment.put(consumer, ownedPartitions);
- // 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<TopicPartition, TreeMap<Integer, String>> partitionConsumersEntry: sortedPartitionConsumersByGeneration.entrySet()) {
- TopicPartition partition = partitionConsumersEntry.getKey();
- TreeMap<Integer, String> consumers = partitionConsumersEntry.getValue();
- Iterator<Integer> it = consumers.descendingKeySet().iterator();
+ // Only consider this consumer's owned partitions as valid if it is a member of the current highest
+ // generation, or it's generation is not present but we have not seen any known generation so far
+ if (memberData.generation.isPresent() && memberData.generation.get() >= maxGeneration
+ || !memberData.generation.isPresent() && maxGeneration == DEFAULT_GENERATION) {
- // 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);
+ // If the current member's generation is higher, all the previously owned partitions are invalid
+ if (memberData.generation.isPresent() && memberData.generation.get() > maxGeneration) {
+ for (String member: membersOfCurrentHighestGeneration) {
+ List<TopicPartition> oldGenerationPartitions = currentAssignment.get(member);
+ updatePrevAssignment(prevAssignment, oldGenerationPartitions, member, maxGeneration);
+ oldGenerationPartitions.clear();
+ }
+ membersOfCurrentHighestGeneration.clear();
+ maxGeneration = memberData.generation.get();
+ }
- // now update previous assignment if any
- if (it.hasNext()) {
- int generation = it.next();
- prevAssignment.put(partition, new ConsumerGenerationPair(consumers.get(generation), generation));
+ membersOfCurrentHighestGeneration.add(consumer);
+ ownedPartitions.addAll(memberData.partitions);
+ } else if (!memberData.generation.isPresent()) {
+ // current maxGeneration is larger than DEFAULT_GENERATION,
+ // put all partitions as DEFAULT_GENERATION into provAssignment
Review comment:
I'm not 100% sure if the refactored code has the exactly same logic as the old code now since its branching conditions have largely changed. E.g. do we still detect if a partition is assigned to different consumers in a generation or not?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org