From 8bd9a40a64cb178aca2720f4161dc1fa90b1b6de Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Mon, 25 Jun 2018 18:38:38 -0700 Subject: [PATCH 01/17] KAFKA-7026: Sticky Assignor Partition Assignment Improvement In the current implementation of sticky assignor the leader does not cache the most recent calculated assignment. It relies on the fact that each consumer in the group sends its subscribed topics and also its current assignment when a rebalance occurs. This could lead to the issue described in KAFKA-7026, in which current assignment of a consumer is no loonger valid and should be ignored. The solution implemented in this PR involves the leader caching the most recent assignment of each consumer, so the assignment reported by a consumer can be properly ignored if necessary. --- .../clients/consumer/StickyAssignor.java | 68 ++++++++++--------- 1 file changed, 36 insertions(+), 32 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index ee537eba78802..a3d8aecde356f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -194,15 +194,14 @@ public class StickyAssignor extends AbstractPartitionAssignor { private static final Schema STICKY_ASSIGNOR_USER_DATA = new Schema( new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT))); + private Map> currentAssignment = new HashMap<>(); private List memberAssignment = null; private PartitionMovements partitionMovements; public Map> assign(Map partitionsPerTopic, Map subscriptions) { - Map> currentAssignment = new HashMap<>(); partitionMovements = new PartitionMovements(); - - prepopulateCurrentAssignments(subscriptions, currentAssignment); + prepopulateCurrentAssignments(subscriptions); boolean isFreshAssignment = currentAssignment.isEmpty(); // a mapping of all topic partitions to all consumers that can be assigned to them @@ -239,7 +238,7 @@ public Map> assign(Map partitionsP currentPartitionConsumer.put(topicPartition, entry.getKey()); List sortedPartitions = sortPartitions( - currentAssignment, isFreshAssignment, partition2AllPotentialConsumers, consumer2AllPotentialPartitions); + isFreshAssignment, partition2AllPotentialConsumers, consumer2AllPotentialPartitions); // all partitions that need to be assigned (initially set to all partitions but adjusted in the following loop) List unassignedPartitions = new ArrayList<>(sortedPartitions); @@ -278,17 +277,36 @@ public Map> assign(Map partitionsP TreeSet sortedCurrentSubscriptions = new TreeSet<>(new SubscriptionComparator(currentAssignment)); sortedCurrentSubscriptions.addAll(currentAssignment.keySet()); - balance(currentAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions, + balance(sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer); return currentAssignment; } - private void prepopulateCurrentAssignments(Map subscriptions, - Map> currentAssignment) { + private void prepopulateCurrentAssignments(Map subscriptions) { + Map currentPartitionToMemberAssignment = new HashMap<>(); + currentAssignment.forEach((member, partitions) -> partitions.stream().forEach( + partition -> currentPartitionToMemberAssignment.put(partition, member) + )); + + currentAssignment.clear(); + mainLoop: for (Map.Entry subscriptionEntry : subscriptions.entrySet()) { ByteBuffer userData = subscriptionEntry.getValue().userData(); - if (userData != null && userData.hasRemaining()) - currentAssignment.put(subscriptionEntry.getKey(), deserializeTopicPartitionAssignment(userData)); + if (userData != null && userData.hasRemaining()) { + String member = subscriptionEntry.getKey(); + List deserializedAssignment = deserializeTopicPartitionAssignment(userData); + for (TopicPartition partition: deserializedAssignment) { + // this check is the reason for having the leader cache the current assignment + if (currentPartitionToMemberAssignment.containsKey(partition) && !currentPartitionToMemberAssignment.get(partition).equals(member)) { + log.debug("Member " + member + " seems to be out-of-date with respect to its previous assignment " + + "(maybe it's rejoining the group after some pause?). Therefore, It's previous assignment " + + "will be ignored."); + continue mainLoop; + } + } + + currentAssignment.put(member, deserializedAssignment); + } } } @@ -313,13 +331,11 @@ public String name() { /** * 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 * @return true if the given assignment is balanced; false otherwise */ - private boolean isBalanced(Map> currentAssignment, - TreeSet sortedCurrentSubscriptions, + private boolean isBalanced(TreeSet sortedCurrentSubscriptions, Map> allSubscriptions) { int min = currentAssignment.get(sortedCurrentSubscriptions.first()).size(); int max = currentAssignment.get(sortedCurrentSubscriptions.last()).size(); @@ -394,14 +410,12 @@ private int getBalanceScore(Map> 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 currentAssignment the calculated assignment so far * @param isFreshAssignment whether this is a new assignment, or a reassignment of an existing one * @param partition2AllPotentialConsumers a mapping of partitions to their potential consumers * @param consumer2AllPotentialPartitions a mapping of consumers to potential partitions they can consumer from * @return sorted list of valid partitions */ - private List sortPartitions(Map> currentAssignment, - boolean isFreshAssignment, + private List sortPartitions(boolean isFreshAssignment, Map> partition2AllPotentialConsumers, Map> consumer2AllPotentialPartitions) { List sortedPartitions = new ArrayList<>(); @@ -471,7 +485,6 @@ private boolean areSubscriptionsIdentical(Map> part */ private String assignPartition(TopicPartition partition, TreeSet sortedCurrentSubscriptions, - Map> currentAssignment, Map> consumer2AllPotentialPartitions, Map currentPartitionConsumer) { for (String consumer: sortedCurrentSubscriptions) { @@ -493,7 +506,6 @@ private boolean canParticipateInReassignment(TopicPartition partition, } private boolean canParticipateInReassignment(String consumer, - Map> currentAssignment, Map> consumer2AllPotentialPartitions, Map> partition2AllPotentialConsumers) { List currentPartitions = currentAssignment.get(consumer); @@ -518,8 +530,7 @@ private boolean canParticipateInReassignment(String consumer, /** * Balance the current assignment using the data structures created in the assign(...) method above. */ - private void balance(Map> currentAssignment, - List sortedPartitions, + private void balance(List sortedPartitions, List unassignedPartitions, TreeSet sortedCurrentSubscriptions, Map> consumer2AllPotentialPartitions, @@ -534,8 +545,7 @@ private void balance(Map> currentAssignment, if (partition2AllPotentialConsumers.get(partition).isEmpty()) continue; - assignPartition(partition, sortedCurrentSubscriptions, currentAssignment, - consumer2AllPotentialPartitions, currentPartitionConsumer); + assignPartition(partition, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, currentPartitionConsumer); } // narrow down the reassignment scope to only those partitions that can actually be reassigned @@ -548,8 +558,7 @@ private void balance(Map> currentAssignment, // narrow down the reassignment scope to only those consumers that are subject to reassignment Map> fixedAssignments = new HashMap<>(); for (String consumer: consumer2AllPotentialPartitions.keySet()) - if (!canParticipateInReassignment(consumer, currentAssignment, - consumer2AllPotentialPartitions, partition2AllPotentialConsumers)) { + if (!canParticipateInReassignment(consumer, consumer2AllPotentialPartitions, partition2AllPotentialConsumers)) { sortedCurrentSubscriptions.remove(consumer); fixedAssignments.put(consumer, currentAssignment.remove(consumer)); } @@ -558,7 +567,7 @@ private void balance(Map> currentAssignment, Map> preBalanceAssignment = deepCopy(currentAssignment); Map preBalancePartitionConsumers = new HashMap<>(currentPartitionConsumer); - reassignmentPerformed = performReassignments(sortedPartitions, currentAssignment, sortedCurrentSubscriptions, + reassignmentPerformed = performReassignments(sortedPartitions, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer); // if we are not preserving existing assignments and we have made changes to the current assignment @@ -580,7 +589,6 @@ private void balance(Map> currentAssignment, } private boolean performReassignments(List reassignablePartitions, - Map> currentAssignment, TreeSet sortedCurrentSubscriptions, Map> consumer2AllPotentialPartitions, Map> partition2AllPotentialConsumers, @@ -594,7 +602,7 @@ 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(sortedCurrentSubscriptions, consumer2AllPotentialPartitions)) { TopicPartition partition = partitionIterator.next(); // the partition must have at least two consumers @@ -609,7 +617,7 @@ 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)) { if (currentAssignment.get(consumer).size() > currentAssignment.get(otherConsumer).size() + 1) { - reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, consumer2AllPotentialPartitions); + reassignPartition(partition, sortedCurrentSubscriptions, currentPartitionConsumer, consumer2AllPotentialPartitions); reassignmentPerformed = true; modified = true; break; @@ -622,7 +630,6 @@ private boolean performReassignments(List reassignablePartitions } private void reassignPartition(TopicPartition partition, - Map> currentAssignment, TreeSet sortedCurrentSubscriptions, Map currentPartitionConsumer, Map> consumer2AllPotentialPartitions) { @@ -641,14 +648,11 @@ private void reassignPartition(TopicPartition partition, // find the correct partition movement considering the stickiness requirement TopicPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(partition, consumer, newConsumer); - processPartitionMovement(partitionToBeMoved, newConsumer, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer); - - return; + processPartitionMovement(partitionToBeMoved, newConsumer, sortedCurrentSubscriptions, currentPartitionConsumer); } private void processPartitionMovement(TopicPartition partition, String newConsumer, - Map> currentAssignment, TreeSet sortedCurrentSubscriptions, Map currentPartitionConsumer) { String oldConsumer = currentPartitionConsumer.get(partition); From 94edbcc59bfadbabb9b784efcf6526593311aa6a Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Tue, 3 Jul 2018 11:03:42 -0700 Subject: [PATCH 02/17] Changed the solution to use generation markers instead of caching the assignment by the leader This helps in situations where a new leader is chosen and there are dupliacates in the reported assignments by consumers. --- .../clients/consumer/StickyAssignor.java | 195 ++++++++++++++---- .../clients/consumer/StickyAssignorTest.java | 174 ++++++++++++++-- 2 files changed, 311 insertions(+), 58 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index a3d8aecde356f..a0ee7b94945d6 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -30,6 +30,7 @@ import java.io.Serializable; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.Comparator; @@ -188,20 +189,45 @@ public class StickyAssignor extends AbstractPartitionAssignor { private static final String TOPIC_PARTITIONS_KEY_NAME = "previous_assignment"; private static final String TOPIC_KEY_NAME = "topic"; private static final String PARTITIONS_KEY_NAME = "partitions"; + private static final String GENERATION_KEY_NAME = "generation"; private static final Schema TOPIC_ASSIGNMENT = new Schema( new Field(TOPIC_KEY_NAME, Type.STRING), new Field(PARTITIONS_KEY_NAME, new ArrayOf(Type.INT32))); private static final Schema STICKY_ASSIGNOR_USER_DATA = new Schema( - new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT))); + new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT)), + new Field(GENERATION_KEY_NAME, Type.INT32)); - private Map> currentAssignment = new HashMap<>(); private List memberAssignment = null; + // this generation is specific to sticky assignor and is used to validate + // assignments reported by each consumer during a rebalance + private int generation = 0; private PartitionMovements partitionMovements; + static final class ConsumerUserData { + List partitions; + int generation; + ConsumerUserData(List partitions, int generation) { + this.partitions = partitions; + this.generation = generation; + } + } + + static final class ConsumerGenerationPair { + String consumer; + int generation; + ConsumerGenerationPair(String consumer, int generation) { + this.consumer = consumer; + this.generation = generation; + } + } + public Map> assign(Map partitionsPerTopic, Map subscriptions) { + Map> currentAssignment = new HashMap<>(); + Map prevAssignment = new HashMap<>(); partitionMovements = new PartitionMovements(); - prepopulateCurrentAssignments(subscriptions); + + prepopulateCurrentAssignments(subscriptions, currentAssignment, prevAssignment); boolean isFreshAssignment = currentAssignment.isEmpty(); // a mapping of all topic partitions to all consumers that can be assigned to them @@ -238,7 +264,7 @@ public Map> assign(Map partitionsP currentPartitionConsumer.put(topicPartition, entry.getKey()); List sortedPartitions = sortPartitions( - isFreshAssignment, partition2AllPotentialConsumers, consumer2AllPotentialPartitions); + currentAssignment, prevAssignment, isFreshAssignment, partition2AllPotentialConsumers, consumer2AllPotentialPartitions); // all partitions that need to be assigned (initially set to all partitions but adjusted in the following loop) List unassignedPartitions = new ArrayList<>(sortedPartitions); @@ -277,37 +303,80 @@ public Map> assign(Map partitionsP TreeSet sortedCurrentSubscriptions = new TreeSet<>(new SubscriptionComparator(currentAssignment)); sortedCurrentSubscriptions.addAll(currentAssignment.keySet()); - balance(sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions, + balance(currentAssignment, prevAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer); return currentAssignment; } - private void prepopulateCurrentAssignments(Map subscriptions) { - Map currentPartitionToMemberAssignment = new HashMap<>(); - currentAssignment.forEach((member, partitions) -> partitions.stream().forEach( - partition -> currentPartitionToMemberAssignment.put(partition, member) - )); + private void updatePrevAssignment(Map prevAssignment, + TopicPartition partition, ConsumerGenerationPair consumerGeneration) { + if (!prevAssignment.containsKey(partition)) + prevAssignment.put(partition, consumerGeneration); + else { + ConsumerGenerationPair existingPair = prevAssignment.get(partition); + if (consumerGeneration.generation > existingPair.generation) + prevAssignment.put(partition, consumerGeneration); + else if (consumerGeneration.generation == existingPair.generation) + // same partition is assigned to two consumers during the same rebalance + throw new IllegalStateException("Error: Partition '" + partition + "' had been assigned to " + + "multiple consumers following sticky assignment generation " + consumerGeneration.generation + "."); + } + } - currentAssignment.clear(); - mainLoop: + 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 + + Map partitionConsumer = new HashMap<>(); + int maxGeneration = 0; for (Map.Entry subscriptionEntry : subscriptions.entrySet()) { + String consumer = subscriptionEntry.getKey(); ByteBuffer userData = subscriptionEntry.getValue().userData(); - if (userData != null && userData.hasRemaining()) { - String member = subscriptionEntry.getKey(); - List deserializedAssignment = deserializeTopicPartitionAssignment(userData); - for (TopicPartition partition: deserializedAssignment) { - // this check is the reason for having the leader cache the current assignment - if (currentPartitionToMemberAssignment.containsKey(partition) && !currentPartitionToMemberAssignment.get(partition).equals(member)) { - log.debug("Member " + member + " seems to be out-of-date with respect to its previous assignment " + - "(maybe it's rejoining the group after some pause?). Therefore, It's previous assignment " + - "will be ignored."); - continue mainLoop; + if (userData == null || !userData.hasRemaining()) continue; + ConsumerUserData consumerUserData = deserializeTopicPartitionAssignment(userData); + if (consumerUserData.generation > maxGeneration) + maxGeneration = consumerUserData.generation; + + for (TopicPartition partition: consumerUserData.partitions) { + if (!partitionConsumer.containsKey(partition)) + partitionConsumer.put(partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); + else { + ConsumerGenerationPair existingRecord = partitionConsumer.get(partition); + if (consumerUserData.generation == existingRecord.generation) { + // same partition is assigned to two consumers during the same rebalance + throw new IllegalStateException("Error: Partition '" + partition + "' is assigned to multiple consumers " + + "following sticky assignment generation " + consumerUserData.generation + "."); + } else if (consumerUserData.generation > existingRecord.generation) { + // same partition is assigned to two consumers in different rebalances. + // the assignment being processed has higher generation that the one already processed, and takes priority + updatePrevAssignment(prevAssignment, partition, existingRecord); + partitionConsumer.put(partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); + } else { + // else if (consumerUserData.generation < existingRecord.generation) + // if the same partition is assigned to two consumers in different rebalances and + // the assignment being processed has lower generation that the one already processed, + // it is only considered for updating the previous assignment of the involved topic partition + updatePrevAssignment(prevAssignment, partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); } } + } + } - currentAssignment.put(member, deserializedAssignment); + for (Map.Entry partitionAssignment: partitionConsumer.entrySet()) { + TopicPartition partition = partitionAssignment.getKey(); + String consumer = partitionAssignment.getValue().consumer; + if (currentAssignment.containsKey(consumer)) + currentAssignment.get(consumer).add(partition); + else { + List partitions = new ArrayList<>(Arrays.asList(partition)); + currentAssignment.put(consumer, partitions); } } + + this.generation = maxGeneration + 1; } @Override @@ -320,7 +389,8 @@ public Subscription subscription(Set topics) { if (memberAssignment == null) return new Subscription(new ArrayList<>(topics)); - return new Subscription(new ArrayList<>(topics), serializeTopicPartitionAssignment(memberAssignment)); + return new Subscription(new ArrayList<>(topics), + serializeTopicPartitionAssignment(new ConsumerUserData(memberAssignment, generation))); } @Override @@ -328,14 +398,20 @@ public String name() { return "sticky"; } + public int generation() { + return generation; + } + /** * 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 * @return true if the given assignment is balanced; false otherwise */ - private boolean isBalanced(TreeSet sortedCurrentSubscriptions, + private boolean isBalanced(Map> currentAssignment, + TreeSet sortedCurrentSubscriptions, Map> allSubscriptions) { int min = currentAssignment.get(sortedCurrentSubscriptions.first()).size(); int max = currentAssignment.get(sortedCurrentSubscriptions.last()).size(); @@ -410,12 +486,16 @@ private int getBalanceScore(Map> 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 currentAssignment the calculated assignment so far + * @param prevAssignment previous consumer of each partition, if any * @param isFreshAssignment whether this is a new assignment, or a reassignment of an existing one * @param partition2AllPotentialConsumers a mapping of partitions to their potential consumers * @param consumer2AllPotentialPartitions a mapping of consumers to potential partitions they can consumer from * @return sorted list of valid partitions */ - private List sortPartitions(boolean isFreshAssignment, + private List sortPartitions(Map> currentAssignment, + Map prevAssignment, + boolean isFreshAssignment, Map> partition2AllPotentialConsumers, Map> consumer2AllPotentialPartitions) { List sortedPartitions = new ArrayList<>(); @@ -436,10 +516,17 @@ private List sortPartitions(boolean isFreshAssignment, TreeSet sortedConsumers = new TreeSet<>(new SubscriptionComparator(assignments)); sortedConsumers.addAll(assignments.keySet()); + List prevPartitions = new ArrayList(prevAssignment.keySet()); while (!sortedConsumers.isEmpty()) { String consumer = sortedConsumers.pollLast(); List remainingPartitions = assignments.get(consumer); - if (!remainingPartitions.isEmpty()) { + prevPartitions.retainAll(remainingPartitions); + if (!prevPartitions.isEmpty()) { + TopicPartition partition = prevPartitions.remove(0); + remainingPartitions.remove(partition); + sortedPartitions.add(partition); + sortedConsumers.add(consumer); + } else if (!remainingPartitions.isEmpty()) { sortedPartitions.add(remainingPartitions.remove(0)); sortedConsumers.add(consumer); } @@ -485,6 +572,7 @@ private boolean areSubscriptionsIdentical(Map> part */ private String assignPartition(TopicPartition partition, TreeSet sortedCurrentSubscriptions, + Map> currentAssignment, Map> consumer2AllPotentialPartitions, Map currentPartitionConsumer) { for (String consumer: sortedCurrentSubscriptions) { @@ -506,6 +594,7 @@ private boolean canParticipateInReassignment(TopicPartition partition, } private boolean canParticipateInReassignment(String consumer, + Map> currentAssignment, Map> consumer2AllPotentialPartitions, Map> partition2AllPotentialConsumers) { List currentPartitions = currentAssignment.get(consumer); @@ -530,7 +619,9 @@ private boolean canParticipateInReassignment(String consumer, /** * Balance the current assignment using the data structures created in the assign(...) method above. */ - private void balance(List sortedPartitions, + private void balance(Map> currentAssignment, + Map prevAssignment, + List sortedPartitions, List unassignedPartitions, TreeSet sortedCurrentSubscriptions, Map> consumer2AllPotentialPartitions, @@ -545,7 +636,8 @@ private void balance(List sortedPartitions, if (partition2AllPotentialConsumers.get(partition).isEmpty()) continue; - assignPartition(partition, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, currentPartitionConsumer); + assignPartition(partition, sortedCurrentSubscriptions, currentAssignment, + consumer2AllPotentialPartitions, currentPartitionConsumer); } // narrow down the reassignment scope to only those partitions that can actually be reassigned @@ -558,7 +650,8 @@ private void balance(List sortedPartitions, // narrow down the reassignment scope to only those consumers that are subject to reassignment Map> fixedAssignments = new HashMap<>(); for (String consumer: consumer2AllPotentialPartitions.keySet()) - if (!canParticipateInReassignment(consumer, consumer2AllPotentialPartitions, partition2AllPotentialConsumers)) { + if (!canParticipateInReassignment(consumer, currentAssignment, + consumer2AllPotentialPartitions, partition2AllPotentialConsumers)) { sortedCurrentSubscriptions.remove(consumer); fixedAssignments.put(consumer, currentAssignment.remove(consumer)); } @@ -567,7 +660,7 @@ private void balance(List sortedPartitions, Map> preBalanceAssignment = deepCopy(currentAssignment); Map preBalancePartitionConsumers = new HashMap<>(currentPartitionConsumer); - reassignmentPerformed = performReassignments(sortedPartitions, sortedCurrentSubscriptions, + reassignmentPerformed = performReassignments(sortedPartitions, currentAssignment, prevAssignment, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer); // if we are not preserving existing assignments and we have made changes to the current assignment @@ -589,6 +682,8 @@ private void balance(List sortedPartitions, } private boolean performReassignments(List reassignablePartitions, + Map> currentAssignment, + Map prevAssignment, TreeSet sortedCurrentSubscriptions, Map> consumer2AllPotentialPartitions, Map> partition2AllPotentialConsumers, @@ -602,7 +697,7 @@ 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(sortedCurrentSubscriptions, consumer2AllPotentialPartitions)) { + while (partitionIterator.hasNext() && !isBalanced(currentAssignment, sortedCurrentSubscriptions, consumer2AllPotentialPartitions)) { TopicPartition partition = partitionIterator.next(); // the partition must have at least two consumers @@ -614,10 +709,18 @@ private boolean performReassignments(List reassignablePartitions 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, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, prevAssignment.get(partition).consumer); + reassignmentPerformed = true; + modified = true; + continue; + } + // check if a better-suited consumer exist for the partition; if so, reassign it for (String otherConsumer: partition2AllPotentialConsumers.get(partition)) { if (currentAssignment.get(consumer).size() > currentAssignment.get(otherConsumer).size() + 1) { - reassignPartition(partition, sortedCurrentSubscriptions, currentPartitionConsumer, consumer2AllPotentialPartitions); + reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, consumer2AllPotentialPartitions); reassignmentPerformed = true; modified = true; break; @@ -630,11 +733,10 @@ private boolean performReassignments(List reassignablePartitions } private void reassignPartition(TopicPartition partition, + Map> currentAssignment, TreeSet sortedCurrentSubscriptions, Map currentPartitionConsumer, Map> consumer2AllPotentialPartitions) { - String consumer = currentPartitionConsumer.get(partition); - // find the new consumer String newConsumer = null; for (String anotherConsumer: sortedCurrentSubscriptions) { @@ -646,13 +748,23 @@ private void reassignPartition(TopicPartition partition, assert newConsumer != null; + reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, newConsumer); + } + + private void reassignPartition(TopicPartition partition, + Map> currentAssignment, + TreeSet sortedCurrentSubscriptions, + Map currentPartitionConsumer, + 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, sortedCurrentSubscriptions, currentPartitionConsumer); + processPartitionMovement(partitionToBeMoved, newConsumer, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer); } private void processPartitionMovement(TopicPartition partition, String newConsumer, + Map> currentAssignment, TreeSet sortedCurrentSubscriptions, Map currentPartitionConsumer) { String oldConsumer = currentPartitionConsumer.get(partition); @@ -673,23 +785,24 @@ boolean isSticky() { return partitionMovements.isSticky(); } - static ByteBuffer serializeTopicPartitionAssignment(List partitions) { + static ByteBuffer serializeTopicPartitionAssignment(ConsumerUserData consumerUserData) { Struct struct = new Struct(STICKY_ASSIGNOR_USER_DATA); List topicAssignments = new ArrayList<>(); - for (Map.Entry> topicEntry : CollectionUtils.groupPartitionsByTopic(partitions).entrySet()) { + for (Map.Entry> topicEntry : CollectionUtils.groupPartitionsByTopic(consumerUserData.partitions).entrySet()) { Struct topicAssignment = new Struct(TOPIC_ASSIGNMENT); topicAssignment.set(TOPIC_KEY_NAME, topicEntry.getKey()); topicAssignment.set(PARTITIONS_KEY_NAME, topicEntry.getValue().toArray()); topicAssignments.add(topicAssignment); } struct.set(TOPIC_PARTITIONS_KEY_NAME, topicAssignments.toArray()); + struct.set(GENERATION_KEY_NAME, consumerUserData.generation); ByteBuffer buffer = ByteBuffer.allocate(STICKY_ASSIGNOR_USER_DATA.sizeOf(struct)); STICKY_ASSIGNOR_USER_DATA.write(buffer, struct); buffer.flip(); return buffer; } - private static List deserializeTopicPartitionAssignment(ByteBuffer buffer) { + private static ConsumerUserData deserializeTopicPartitionAssignment(ByteBuffer buffer) { Struct struct = STICKY_ASSIGNOR_USER_DATA.read(buffer); List partitions = new ArrayList<>(); for (Object structObj : struct.getArray(TOPIC_PARTITIONS_KEY_NAME)) { @@ -700,7 +813,9 @@ private static List deserializeTopicPartitionAssignment(ByteBuff partitions.add(new TopicPartition(topic, partition)); } } - return partitions; + // make sure this is backward compatible + int generation = struct.hasField(GENERATION_KEY_NAME) ? struct.getInt(GENERATION_KEY_NAME) : 0; + return new ConsumerUserData(partitions, generation); } /** 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 32ba16a482023..be2651d1b095c 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 @@ -30,6 +30,7 @@ import java.util.Random; import java.util.Set; +import org.apache.kafka.clients.consumer.StickyAssignor.ConsumerUserData; import org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.utils.CollectionUtils; @@ -235,10 +236,11 @@ public void testAddRemoveConsumerOneTopic() { String consumer2 = "consumer2"; subscriptions.put(consumer1, - new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer1)))); + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(assignment.get(consumer1), assignor.generation())))); subscriptions.put(consumer2, new Subscription(topics(topic))); assignment = assignor.assign(partitionsPerTopic, subscriptions); - assertEquals(partitions(tp(topic, 1), tp(topic, 2)), assignment.get(consumer1)); + assertEquals(partitions(tp(topic, 2), tp(topic, 1)), assignment.get(consumer1)); assertEquals(partitions(tp(topic, 0)), assignment.get(consumer2)); verifyValidityAndBalance(subscriptions, assignment); @@ -247,7 +249,8 @@ public void testAddRemoveConsumerOneTopic() { subscriptions.remove(consumer1); subscriptions.put(consumer2, - new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer2)))); + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(assignment.get(consumer2), assignor.generation())))); assignment = assignor.assign(partitionsPerTopic, subscriptions); assertTrue(assignment.get(consumer2).contains(tp(topic, 0))); assertTrue(assignment.get(consumer2).contains(tp(topic, 1))); @@ -318,9 +321,11 @@ public void testAddRemoveTopicTwoConsumers() { String topic2 = "topic2"; partitionsPerTopic.put(topic2, 3); subscriptions.put(consumer1, - new Subscription(topics(topic, topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer1)))); + new Subscription(topics(topic, topic2), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(assignment.get(consumer1), assignor.generation())))); subscriptions.put(consumer2, - new Subscription(topics(topic, topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer2)))); + new Subscription(topics(topic, topic2), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(assignment.get(consumer2), assignor.generation())))); assignment = assignor.assign(partitionsPerTopic, subscriptions); // verify balance verifyValidityAndBalance(subscriptions, assignment); @@ -335,9 +340,11 @@ public void testAddRemoveTopicTwoConsumers() { partitionsPerTopic.remove(topic); subscriptions.put(consumer1, - new Subscription(topics(topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer1)))); + new Subscription(topics(topic2), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(assignment.get(consumer1), assignor.generation())))); subscriptions.put(consumer2, - new Subscription(topics(topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer2)))); + new Subscription(topics(topic2), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(assignment.get(consumer2), assignor.generation())))); assignment = assignor.assign(partitionsPerTopic, subscriptions); // verify balance verifyValidityAndBalance(subscriptions, assignment); @@ -372,7 +379,8 @@ public void testReassignmentAfterOneConsumerLeaves() { for (int i = 1; i < 20; i++) { String consumer = getConsumerName(i, 20); subscriptions.put(consumer, - new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer)))); + new Subscription(subscriptions.get(consumer).topics(), + StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), assignor.generation())))); } subscriptions.remove("consumer10"); @@ -421,7 +429,8 @@ public void testSameSubscriptions() { for (int i = 1; i < 9; i++) { String consumer = getConsumerName(i, 9); subscriptions.put(consumer, - new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer)))); + new Subscription(subscriptions.get(consumer).topics(), + StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), assignor.generation())))); } subscriptions.remove(getConsumerName(5, 9)); @@ -454,7 +463,8 @@ public void testLargeAssignmentWithMultipleConsumersLeaving() { for (int i = 1; i < consumerCount; i++) { String consumer = getConsumerName(i, consumerCount); subscriptions.put(consumer, - new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer)))); + new Subscription(subscriptions.get(consumer).topics(), + StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), assignor.generation())))); } for (int i = 0; i < 50; ++i) { String c = getConsumerName(rand.nextInt(consumerCount), consumerCount); @@ -526,7 +536,7 @@ public void testReassignmentWithRandomSubscriptionsAndChanges() { List sub = Utils.sorted(getRandomSublist(topics)); String consumer = getConsumerName(i, maxNumConsumers); subscriptions.put(consumer, - new Subscription(sub, StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer)))); + new Subscription(sub, StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), assignor.generation())))); } assignment = assignor.assign(partitionsPerTopic, subscriptions); @@ -544,13 +554,16 @@ public void testMoveExistingAssignments() { Map subscriptions = new HashMap<>(); subscriptions.put("consumer01", new Subscription(topics("topic01", "topic02"), - StickyAssignor.serializeTopicPartitionAssignment(partitions(tp("topic01", 0))))); + StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(partitions(tp("topic01", 0)), assignor.generation())))); subscriptions.put("consumer02", new Subscription(topics("topic01", "topic02", "topic03", "topic04"), - StickyAssignor.serializeTopicPartitionAssignment(partitions(tp("topic02", 0), tp("topic03", 0))))); + StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(partitions(tp("topic02", 0), tp("topic03", 0)), assignor.generation())))); subscriptions.put("consumer03", new Subscription(topics("topic02", "topic03", "topic04", "topic05", "topic06"), - StickyAssignor.serializeTopicPartitionAssignment(partitions(tp("topic04", 0), tp("topic05", 0), tp("topic06", 0))))); + StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(partitions(tp("topic04", 0), tp("topic05", 0), tp("topic06", 0)), assignor.generation())))); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment); @@ -584,13 +597,16 @@ public void testStickiness() { subscriptions.remove("consumer01"); subscriptions.put("consumer02", new Subscription(topics("topic01"), - StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer02")))); + StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(assignment.get("consumer02"), assignor.generation())))); subscriptions.put("consumer03", new Subscription(topics("topic01"), - StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer03")))); + StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(assignment.get("consumer03"), assignor.generation())))); subscriptions.put("consumer04", new Subscription(topics("topic01"), - StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer04")))); + StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(assignment.get("consumer04"), assignor.generation())))); assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment); @@ -631,13 +647,135 @@ public void testNoExceptionThrownWhenOnlySubscribedTopicDeleted() { Map subscriptions = new HashMap<>(); subscriptions.put(consumer, new Subscription(topics(topic))); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); - subscriptions.put(consumer, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer)))); + subscriptions.put(consumer, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(assignment.get(consumer), 1)))); assignment = assignor.assign(Collections.emptyMap(), subscriptions); assertEquals(assignment.size(), 1); assertTrue(assignment.get(consumer).isEmpty()); } + @Test + public void testAssignmentWithMultipleGenerations1() { + String topic = "topic"; + String consumer1 = "consumer1"; + String consumer2 = "consumer2"; + String consumer3 = "consumer3"; + + Map partitionsPerTopic = new HashMap<>(); + partitionsPerTopic.put(topic, 6); + Map subscriptions = new HashMap<>(); + subscriptions.put(consumer1, new Subscription(topics(topic))); + subscriptions.put(consumer2, new Subscription(topics(topic))); + subscriptions.put(consumer3, new Subscription(topics(topic))); + + 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); + assertEquals(1, assignor.generation()); + verifyValidityAndBalance(subscriptions, assignment); + assertTrue(isFullyBalanced(assignment)); + + subscriptions.put(consumer1, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(r1partitions1, 1)))); + subscriptions.put(consumer2, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(r1partitions2, 1)))); + 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); + assertEquals(2, assignor.generation()); + assertTrue(r2partitions1.containsAll(r1partitions1)); + assertTrue(r2partitions2.containsAll(r1partitions2)); + verifyValidityAndBalance(subscriptions, assignment); + assertTrue(isFullyBalanced(assignment)); + assertTrue(assignor.isSticky()); + + assertTrue(!Collections.disjoint(r2partitions2, r1partitions3)); + subscriptions.remove(consumer1); + subscriptions.put(consumer2, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(r2partitions2, 2)))); + subscriptions.put(consumer3, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(r1partitions3, 1)))); + assignment = assignor.assign(partitionsPerTopic, subscriptions); + List r3partitions2 = assignment.get(consumer2); + List r3partitions3 = assignment.get(consumer3); + assertTrue(r3partitions2.size() == 3 && r3partitions3.size() == 3); + assertEquals(3, assignor.generation()); + assertTrue(Collections.disjoint(r3partitions2, r3partitions3)); + verifyValidityAndBalance(subscriptions, assignment); + assertTrue(isFullyBalanced(assignment)); + assertTrue(assignor.isSticky()); + } + + @Test + public void testAssignmentWithMultipleGenerations2() { + String topic = "topic"; + String consumer1 = "consumer1"; + String consumer2 = "consumer2"; + String consumer3 = "consumer3"; + + Map partitionsPerTopic = new HashMap<>(); + partitionsPerTopic.put(topic, 6); + Map subscriptions = new HashMap<>(); + subscriptions.put(consumer1, new Subscription(topics(topic))); + subscriptions.put(consumer2, new Subscription(topics(topic))); + subscriptions.put(consumer3, new Subscription(topics(topic))); + + 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); + assertEquals(1, assignor.generation()); + verifyValidityAndBalance(subscriptions, assignment); + assertTrue(isFullyBalanced(assignment)); + + subscriptions.remove(consumer1); + subscriptions.put(consumer2, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(r1partitions2, 1)))); + subscriptions.remove(consumer3); + assignment = assignor.assign(partitionsPerTopic, subscriptions); + List r2partitions2 = assignment.get(consumer2); + assertTrue(r2partitions2.size() == 6); + assertEquals(2, assignor.generation()); + assertTrue(r2partitions2.containsAll(r1partitions2)); + verifyValidityAndBalance(subscriptions, assignment); + assertTrue(isFullyBalanced(assignment)); + assertTrue(assignor.isSticky()); + + subscriptions.put(consumer1, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(r1partitions1, 1)))); + subscriptions.put(consumer2, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(r2partitions2, 2)))); + subscriptions.put(consumer3, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(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(3, assignor.generation()); + assertTrue(r3partitions1.equals(r1partitions1)); + assertTrue(r3partitions2.equals(r1partitions2)); + assertTrue(r3partitions3.equals(r1partitions3)); + verifyValidityAndBalance(subscriptions, assignment); + assertTrue(isFullyBalanced(assignment)); + assertTrue(assignor.isSticky()); + } + private String getTopicName(int i, int maxNum) { return getCanonicalName("t", i, maxNum); } From 6ab9400b020f932c4bbc055ac9f8c755b6733f57 Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Fri, 6 Jul 2018 14:47:54 -0700 Subject: [PATCH 03/17] Minor fixes --- .../kafka/clients/consumer/StickyAssignor.java | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index a0ee7b94945d6..c96b72498236b 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -331,7 +331,7 @@ private void prepopulateCurrentAssignments(Map subscriptio // note that a conflict could exists only if user data is for different generations Map partitionConsumer = new HashMap<>(); - int maxGeneration = 0; + int maxGeneration = this.generation; for (Map.Entry subscriptionEntry : subscriptions.entrySet()) { String consumer = subscriptionEntry.getKey(); ByteBuffer userData = subscriptionEntry.getValue().userData(); @@ -346,9 +346,10 @@ private void prepopulateCurrentAssignments(Map subscriptio else { ConsumerGenerationPair existingRecord = partitionConsumer.get(partition); if (consumerUserData.generation == existingRecord.generation) { - // same partition is assigned to two consumers during the same rebalance - throw new IllegalStateException("Error: Partition '" + partition + "' is assigned to multiple consumers " + - "following sticky assignment generation " + consumerUserData.generation + "."); + // same partition is assigned to two consumers during the same rebalance. + // log an error and skip this record + log.error("Error: Partition '" + partition + "' is assigned to multiple consumers " + + "following sticky assignment generation " + consumerUserData.generation + "."); } else if (consumerUserData.generation > existingRecord.generation) { // same partition is assigned to two consumers in different rebalances. // the assignment being processed has higher generation that the one already processed, and takes priority @@ -516,10 +517,11 @@ private List sortPartitions(Map> cu TreeSet sortedConsumers = new TreeSet<>(new SubscriptionComparator(assignments)); sortedConsumers.addAll(assignments.keySet()); - List prevPartitions = new ArrayList(prevAssignment.keySet()); + List prevAssignedPartitions = new ArrayList(prevAssignment.keySet()); while (!sortedConsumers.isEmpty()) { String consumer = sortedConsumers.pollLast(); List remainingPartitions = assignments.get(consumer); + List prevPartitions = new ArrayList<>(prevAssignedPartitions); prevPartitions.retainAll(remainingPartitions); if (!prevPartitions.isEmpty()) { TopicPartition partition = prevPartitions.remove(0); From c0926dbd249a2877ef18943fc9d079d9e26b4333 Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Mon, 9 Jul 2018 14:02:35 -0700 Subject: [PATCH 04/17] Minor update --- .../clients/consumer/StickyAssignor.java | 12 +++-- .../clients/consumer/StickyAssignorTest.java | 48 +++++++++++++++++++ 2 files changed, 55 insertions(+), 5 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index c96b72498236b..c71152ff4c027 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -316,10 +316,12 @@ private void updatePrevAssignment(Map pr ConsumerGenerationPair existingPair = prevAssignment.get(partition); if (consumerGeneration.generation > existingPair.generation) prevAssignment.put(partition, consumerGeneration); - else if (consumerGeneration.generation == existingPair.generation) - // same partition is assigned to two consumers during the same rebalance - throw new IllegalStateException("Error: Partition '" + partition + "' had been assigned to " + - "multiple consumers following sticky assignment generation " + consumerGeneration.generation + "."); + else if (consumerGeneration.generation == existingPair.generation) { + // same partition is assigned to two consumers during the same rebalance. + // log an error and skip this record + log.error("Partition '" + partition + "' had been assigned to multiple consumers " + + "following sticky assignment generation " + consumerGeneration.generation + "."); + } } } @@ -348,7 +350,7 @@ private void prepopulateCurrentAssignments(Map subscriptio if (consumerUserData.generation == existingRecord.generation) { // same partition is assigned to two consumers during the same rebalance. // log an error and skip this record - log.error("Error: Partition '" + partition + "' is assigned to multiple consumers " + + log.error("Partition '" + partition + "' is assigned to multiple consumers " + "following sticky assignment generation " + consumerUserData.generation + "."); } else if (consumerUserData.generation > existingRecord.generation) { // same partition is assigned to two consumers in different rebalances. 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 be2651d1b095c..034bd9dd2d9d9 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 @@ -776,6 +776,54 @@ public void testAssignmentWithMultipleGenerations2() { assertTrue(assignor.isSticky()); } + @Test + public void testAssignmentWithConflictingPreviousGenerations() { + String topic = "topic"; + String consumer1 = "consumer1"; + String consumer2 = "consumer2"; + String consumer3 = "consumer3"; + + Map partitionsPerTopic = new HashMap<>(); + partitionsPerTopic.put(topic, 6); + Map subscriptions = new HashMap<>(); + subscriptions.put(consumer1, new Subscription(topics(topic))); + subscriptions.put(consumer2, new Subscription(topics(topic))); + subscriptions.put(consumer3, new Subscription(topics(topic))); + + 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, tp2, tp3); + List c3partitions0 = partitions(tp3, tp4, tp5); + subscriptions.put(consumer1, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(c1partitions0, 1)))); + subscriptions.put(consumer2, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(c2partitions0, 1)))); + subscriptions.put(consumer3, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(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); + assertEquals(3, assignor.generation()); + assertTrue(c1partitions0.containsAll(c1partitions)); + assertTrue(c2partitions0.containsAll(c2partitions)); + assertTrue(c3partitions0.containsAll(c3partitions)); + verifyValidityAndBalance(subscriptions, assignment); + assertTrue(isFullyBalanced(assignment)); + assertTrue(assignor.isSticky()); + } + private String getTopicName(int i, int maxNum) { return getCanonicalName("t", i, maxNum); } From bde32ed2c9f16e58f7b5a8569809a69572fdec1d Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Wed, 11 Jul 2018 19:07:31 -0700 Subject: [PATCH 05/17] Fixed backward schema compatibility --- .../clients/consumer/StickyAssignor.java | 35 ++++++++---- .../clients/consumer/StickyAssignorTest.java | 57 +++++++++++++++++++ 2 files changed, 82 insertions(+), 10 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index c71152ff4c027..01e43202d2b69 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -186,14 +186,16 @@ public class StickyAssignor extends AbstractPartitionAssignor { // these schemas are used for preserving consumer's previously assigned partitions // list and sending it as user data to the leader during a rebalance - private static final String TOPIC_PARTITIONS_KEY_NAME = "previous_assignment"; - private static final String TOPIC_KEY_NAME = "topic"; - private static final String PARTITIONS_KEY_NAME = "partitions"; + static final String TOPIC_PARTITIONS_KEY_NAME = "previous_assignment"; + static final String TOPIC_KEY_NAME = "topic"; + static final String PARTITIONS_KEY_NAME = "partitions"; private static final String GENERATION_KEY_NAME = "generation"; - private static final Schema TOPIC_ASSIGNMENT = new Schema( + static final Schema TOPIC_ASSIGNMENT = new Schema( new Field(TOPIC_KEY_NAME, Type.STRING), new Field(PARTITIONS_KEY_NAME, new ArrayOf(Type.INT32))); - private static final Schema STICKY_ASSIGNOR_USER_DATA = new Schema( + static final Schema STICKY_ASSIGNOR_USER_DATA_V0 = new Schema( + new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT))); + private static final Schema STICKY_ASSIGNOR_USER_DATA_V1 = new Schema( new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT)), new Field(GENERATION_KEY_NAME, Type.INT32)); @@ -790,7 +792,7 @@ boolean isSticky() { } static ByteBuffer serializeTopicPartitionAssignment(ConsumerUserData consumerUserData) { - Struct struct = new Struct(STICKY_ASSIGNOR_USER_DATA); + Struct struct = new Struct(STICKY_ASSIGNOR_USER_DATA_V1); List topicAssignments = new ArrayList<>(); for (Map.Entry> topicEntry : CollectionUtils.groupPartitionsByTopic(consumerUserData.partitions).entrySet()) { Struct topicAssignment = new Struct(TOPIC_ASSIGNMENT); @@ -800,14 +802,27 @@ static ByteBuffer serializeTopicPartitionAssignment(ConsumerUserData consumerUse } struct.set(TOPIC_PARTITIONS_KEY_NAME, topicAssignments.toArray()); struct.set(GENERATION_KEY_NAME, consumerUserData.generation); - ByteBuffer buffer = ByteBuffer.allocate(STICKY_ASSIGNOR_USER_DATA.sizeOf(struct)); - STICKY_ASSIGNOR_USER_DATA.write(buffer, struct); + ByteBuffer buffer = ByteBuffer.allocate(STICKY_ASSIGNOR_USER_DATA_V1.sizeOf(struct)); + STICKY_ASSIGNOR_USER_DATA_V1.write(buffer, struct); buffer.flip(); return buffer; } private static ConsumerUserData deserializeTopicPartitionAssignment(ByteBuffer buffer) { - Struct struct = STICKY_ASSIGNOR_USER_DATA.read(buffer); + Struct struct; + ByteBuffer copy = buffer.duplicate(); + try { + struct = STICKY_ASSIGNOR_USER_DATA_V1.read(buffer); + } catch (Exception e1) { + try { + // fall back to older schema + struct = STICKY_ASSIGNOR_USER_DATA_V0.read(copy); + } catch (Exception e2) { + // ignore the consumer's previous assignment if it cannot be parsed + return new ConsumerUserData(Arrays.asList(), -2); + } + } + List partitions = new ArrayList<>(); for (Object structObj : struct.getArray(TOPIC_PARTITIONS_KEY_NAME)) { Struct assignment = (Struct) structObj; @@ -818,7 +833,7 @@ private static ConsumerUserData deserializeTopicPartitionAssignment(ByteBuffer b } } // make sure this is backward compatible - int generation = struct.hasField(GENERATION_KEY_NAME) ? struct.getInt(GENERATION_KEY_NAME) : 0; + int generation = struct.hasField(GENERATION_KEY_NAME) ? struct.getInt(GENERATION_KEY_NAME) : -1; return new ConsumerUserData(partitions, generation); } 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 034bd9dd2d9d9..5ac1f3b66be25 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 @@ -19,6 +19,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; +import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -33,6 +34,7 @@ import org.apache.kafka.clients.consumer.StickyAssignor.ConsumerUserData; import org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription; import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.protocol.types.Struct; import org.apache.kafka.common.utils.CollectionUtils; import org.apache.kafka.common.utils.Utils; import org.junit.Test; @@ -824,6 +826,61 @@ public void testAssignmentWithConflictingPreviousGenerations() { assertTrue(assignor.isSticky()); } + @Test + public void testSchemaBackwardCompatibility() { + String topic = "topic"; + String consumer1 = "consumer1"; + String consumer2 = "consumer2"; + String consumer3 = "consumer3"; + + Map partitionsPerTopic = new HashMap<>(); + partitionsPerTopic.put(topic, 3); + Map subscriptions = new HashMap<>(); + subscriptions.put(consumer1, new Subscription(topics(topic))); + subscriptions.put(consumer2, new Subscription(topics(topic))); + subscriptions.put(consumer3, new Subscription(topics(topic))); + + TopicPartition tp0 = new TopicPartition(topic, 0); + TopicPartition tp1 = new TopicPartition(topic, 1); + TopicPartition tp2 = new TopicPartition(topic, 2); + + List c1partitions0 = partitions(tp0, tp2); + List c2partitions0 = partitions(tp1); + subscriptions.put(consumer1, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(c1partitions0, 1)))); + subscriptions.put(consumer2, + new Subscription(topics(topic), serializeTopicPartitionAssignmentToOldSchema(c2partitions0))); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); + List c1partitions = assignment.get(consumer1); + List c2partitions = assignment.get(consumer2); + List c3partitions = assignment.get(consumer3); + + assertTrue(c1partitions.size() == 1 && c2partitions.size() == 1 && c3partitions.size() == 1); + assertEquals(2, assignor.generation()); + assertTrue(c1partitions0.containsAll(c1partitions)); + assertTrue(c2partitions0.containsAll(c2partitions)); + verifyValidityAndBalance(subscriptions, assignment); + assertTrue(isFullyBalanced(assignment)); + assertTrue(assignor.isSticky()); + } + + static ByteBuffer serializeTopicPartitionAssignmentToOldSchema(List partitions) { + Struct struct = new Struct(StickyAssignor.STICKY_ASSIGNOR_USER_DATA_V0); + List topicAssignments = new ArrayList<>(); + for (Map.Entry> topicEntry : CollectionUtils.groupDataByTopic(partitions).entrySet()) { + Struct topicAssignment = new Struct(StickyAssignor.TOPIC_ASSIGNMENT); + topicAssignment.set(StickyAssignor.TOPIC_KEY_NAME, topicEntry.getKey()); + topicAssignment.set(StickyAssignor.PARTITIONS_KEY_NAME, topicEntry.getValue().toArray()); + topicAssignments.add(topicAssignment); + } + struct.set(StickyAssignor.TOPIC_PARTITIONS_KEY_NAME, topicAssignments.toArray()); + ByteBuffer buffer = ByteBuffer.allocate(StickyAssignor.STICKY_ASSIGNOR_USER_DATA_V0.sizeOf(struct)); + StickyAssignor.STICKY_ASSIGNOR_USER_DATA_V0.write(buffer, struct); + buffer.flip(); + return buffer; + } + private String getTopicName(int i, int maxNum) { return getCanonicalName("t", i, maxNum); } From 20c8253133b5a05bd0a80ab603516f9c139be99d Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Tue, 17 Jul 2018 14:04:29 -0700 Subject: [PATCH 06/17] Added a unit test to cover conflicting previous assignments --- .../clients/consumer/StickyAssignorTest.java | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) 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 5ac1f3b66be25..d816a3df94c66 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 @@ -865,6 +865,44 @@ public void testSchemaBackwardCompatibility() { assertTrue(assignor.isSticky()); } + @Test + public void testConflictingPreviousAssignments() { + String topic = "topic"; + String consumer1 = "consumer1"; + String consumer2 = "consumer2"; + + Map partitionsPerTopic = new HashMap<>(); + partitionsPerTopic.put(topic, 2); + Map subscriptions = new HashMap<>(); + subscriptions.put(consumer1, new Subscription(topics(topic))); + subscriptions.put(consumer2, new Subscription(topics(topic))); + + TopicPartition tp0 = new TopicPartition(topic, 0); + TopicPartition tp1 = new TopicPartition(topic, 1); + + // both c1 and c2 have partition 1 assigned to them in generation 1 + List c1partitions0 = partitions(tp0, tp1); + List c2partitions0 = partitions(tp0, tp1); + subscriptions.put(consumer1, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(c1partitions0, 1)))); + subscriptions.put(consumer2, + new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( + new ConsumerUserData(c2partitions0, 1)))); + + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); + List c1partitions = assignment.get(consumer1); + List c2partitions = assignment.get(consumer2); + + assertTrue(c1partitions.size() == 1 && c2partitions.size() == 1); + System.out.println(c1partitions); + System.out.println(c2partitions); + assertEquals(2, assignor.generation()); + verifyValidityAndBalance(subscriptions, assignment); + assertTrue(isFullyBalanced(assignment)); + assertTrue(assignor.isSticky()); + } + static ByteBuffer serializeTopicPartitionAssignmentToOldSchema(List partitions) { Struct struct = new Struct(StickyAssignor.STICKY_ASSIGNOR_USER_DATA_V0); List topicAssignments = new ArrayList<>(); From 7fd53337c8720ddc806d6ea915087733a7c38973 Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Wed, 18 Jul 2018 16:52:45 -0700 Subject: [PATCH 07/17] Switch to using consumer group generation in user data (instead of the locally managed generation) --- .../clients/consumer/StickyAssignor.java | 14 +++++------ .../internals/AbstractCoordinator.java | 4 ++++ .../internals/AbstractPartitionAssignor.java | 24 +++++++++++++++---- .../internals/ConsumerCoordinator.java | 2 +- .../consumer/internals/ConsumerProtocol.java | 1 + .../consumer/internals/PartitionAssignor.java | 21 ++++++++++++---- .../clients/consumer/StickyAssignorTest.java | 20 +++++++--------- 7 files changed, 57 insertions(+), 29 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index 01e43202d2b69..a36eb17c96a7d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -200,10 +200,8 @@ public class StickyAssignor extends AbstractPartitionAssignor { new Field(GENERATION_KEY_NAME, Type.INT32)); private List memberAssignment = null; - // this generation is specific to sticky assignor and is used to validate - // assignments reported by each consumer during a rebalance - private int generation = 0; private PartitionMovements partitionMovements; + private int generation; // consumer group generation static final class ConsumerUserData { List partitions; @@ -223,8 +221,11 @@ static final class ConsumerGenerationPair { } } + @Override public Map> assign(Map partitionsPerTopic, - Map subscriptions) { + Map subscriptions, + int generation) { + this.generation = generation; Map> currentAssignment = new HashMap<>(); Map prevAssignment = new HashMap<>(); partitionMovements = new PartitionMovements(); @@ -335,14 +336,11 @@ private void prepopulateCurrentAssignments(Map subscriptio // note that a conflict could exists only if user data is for different generations Map partitionConsumer = new HashMap<>(); - int maxGeneration = this.generation; for (Map.Entry subscriptionEntry : subscriptions.entrySet()) { String consumer = subscriptionEntry.getKey(); ByteBuffer userData = subscriptionEntry.getValue().userData(); if (userData == null || !userData.hasRemaining()) continue; ConsumerUserData consumerUserData = deserializeTopicPartitionAssignment(userData); - if (consumerUserData.generation > maxGeneration) - maxGeneration = consumerUserData.generation; for (TopicPartition partition: consumerUserData.partitions) { if (!partitionConsumer.containsKey(partition)) @@ -381,7 +379,7 @@ private void prepopulateCurrentAssignments(Map subscriptio } } - this.generation = maxGeneration + 1; + ++generation; } @Override diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java index 4ff4e19590aed..4f29614943c22 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java @@ -753,6 +753,10 @@ protected synchronized Generation generation() { return generation; } + protected Generation latestGeneration() { + return generation; + } + /** * Check whether given generation id is matching the record within current generation. * Only using in unit tests. diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java index 35eb8eb141f93..f7a06903f7761 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java @@ -42,16 +42,30 @@ public abstract class AbstractPartitionAssignor implements PartitionAssignor { * @param subscriptions Map from the memberId to their respective topic subscription * @return Map from each member to the list of partitions assigned to them. */ - public abstract Map> assign(Map partitionsPerTopic, - Map subscriptions); + public Map> assign(Map partitionsPerTopic, + Map subscriptions) { + return assign(partitionsPerTopic, subscriptions, 0); + } - @Override + /** + * Perform the group assignment given the partition counts, member subscriptions, and consumer group generation + * @param partitionsPerTopic The number of partitions for each subscribed topic. Topics not in metadata will be excluded + * from this map. + * @param subscriptions Map from the memberId to their respective topic subscription + * @param generation Consumer group generation id + * @return Map from each member to the list of partitions assigned to them. + */ + public Map> assign(Map partitionsPerTopic, + Map subscriptions, + int generation) { + return assign(partitionsPerTopic, subscriptions); + } @Override public Subscription subscription(Set topics) { return new Subscription(new ArrayList<>(topics)); } @Override - public Map assign(Cluster metadata, Map subscriptions) { + public Map assign(Cluster metadata, Map subscriptions, int generation) { Set allSubscribedTopics = new HashSet<>(); for (Map.Entry subscriptionEntry : subscriptions.entrySet()) allSubscribedTopics.addAll(subscriptionEntry.getValue().topics()); @@ -65,7 +79,7 @@ public Map assign(Cluster metadata, Map> rawAssignments = assign(partitionsPerTopic, subscriptions); + Map> rawAssignments = assign(partitionsPerTopic, subscriptions, generation); // this class maintains no user data, so just wrap the results Map assignments = new HashMap<>(); diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java index 2b949a37deb21..b3964bbd10eaf 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java @@ -411,7 +411,7 @@ protected Map performAssignment(String leaderId, log.debug("Performing assignment using strategy {} with subscriptions {}", assignor.name(), subscriptions); - Map assignment = assignor.assign(metadata.fetch(), subscriptions); + Map assignment = assignor.assign(metadata.fetch(), subscriptions, latestGeneration().generationId); // user-customized assignor may have created some topics that are not in the subscription list // and assign their partitions to the members; in this case we would like to update the leader's diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java index 8a4aef8b33487..7bef8f7ff748a 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java @@ -45,6 +45,7 @@ * TopicPartitions => [Topic Partitions] * Topic => String * Partitions => [int32] + * UserData => Bytes * * * The current implementation assumes that future versions will not break compatibility. When diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java index 4a7c7a8bbd341..93d43a104205d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java @@ -29,7 +29,7 @@ * {@link org.apache.kafka.clients.consumer.KafkaConsumer}. Members of the consumer group subscribe * to the topics they are interested in and forward their subscriptions to a Kafka broker serving * as the group coordinator. The coordinator selects one member to perform the group assignment and - * propagates the subscriptions of all members to it. Then {@link #assign(Cluster, Map)} is called + * propagates the subscriptions of all members to it. Then {@link #assign(Cluster, Map, int)} is called * to perform the assignment and the results are forwarded back to each respective members * * In some cases, it is useful to forward additional metadata to the assignor in order to make @@ -42,7 +42,7 @@ public interface PartitionAssignor { /** * Return a serializable object representing the local member's subscription. This can include * additional information as well (e.g. local host/rack information) which can be leveraged in - * {@link #assign(Cluster, Map)}. + * {@link #assign(Cluster, Map, int)}. * @param topics Topics subscribed to through {@link org.apache.kafka.clients.consumer.KafkaConsumer#subscribe(java.util.Collection)} * and variants * @return Non-null subscription with optional user data @@ -56,12 +56,25 @@ public interface PartitionAssignor { * @return A map from the members to their respective assignment. This should have one entry * for all members who in the input subscription map. */ - Map assign(Cluster metadata, Map subscriptions); + default Map assign(Cluster metadata, Map subscriptions) { + return assign(metadata, subscriptions, 0); + } + /** + * Perform the group assignment given the member subscriptions, current cluster metadata and group generation. + * @param metadata Current topic/broker metadata known by consumer + * @param subscriptions Subscriptions from all members provided through {@link #subscription(Set)} + * @param generation Current group generation id + * @return A map from the members to their respective assignment. This should have one entry + * for all members who in the input subscription map. + */ + default Map assign(Cluster metadata, Map subscriptions, int generation) { + return assign(metadata, subscriptions); + } /** * Callback which is invoked when a group member receives its assignment from the leader. - * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map)} + * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map, int)} */ void onAssignment(Assignment assignment); 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 d816a3df94c66..e4253e5854158 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 @@ -672,7 +672,7 @@ public void testAssignmentWithMultipleGenerations1() { subscriptions.put(consumer2, new Subscription(topics(topic))); subscriptions.put(consumer3, new Subscription(topics(topic))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 0); List r1partitions1 = assignment.get(consumer1); List r1partitions2 = assignment.get(consumer2); List r1partitions3 = assignment.get(consumer3); @@ -688,7 +688,7 @@ public void testAssignmentWithMultipleGenerations1() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions2, 1)))); subscriptions.remove(consumer3); - assignment = assignor.assign(partitionsPerTopic, subscriptions); + assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); List r2partitions1 = assignment.get(consumer1); List r2partitions2 = assignment.get(consumer2); assertTrue(r2partitions1.size() == 3 && r2partitions2.size() == 3); @@ -707,7 +707,7 @@ public void testAssignmentWithMultipleGenerations1() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions3, 1)))); - assignment = assignor.assign(partitionsPerTopic, subscriptions); + assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); List r3partitions2 = assignment.get(consumer2); List r3partitions3 = assignment.get(consumer3); assertTrue(r3partitions2.size() == 3 && r3partitions3.size() == 3); @@ -732,7 +732,7 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.put(consumer2, new Subscription(topics(topic))); subscriptions.put(consumer3, new Subscription(topics(topic))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 0); List r1partitions1 = assignment.get(consumer1); List r1partitions2 = assignment.get(consumer2); List r1partitions3 = assignment.get(consumer3); @@ -746,7 +746,7 @@ public void testAssignmentWithMultipleGenerations2() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions2, 1)))); subscriptions.remove(consumer3); - assignment = assignor.assign(partitionsPerTopic, subscriptions); + assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); List r2partitions2 = assignment.get(consumer2); assertTrue(r2partitions2.size() == 6); assertEquals(2, assignor.generation()); @@ -764,7 +764,7 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions3, 1)))); - assignment = assignor.assign(partitionsPerTopic, subscriptions); + assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); List r3partitions1 = assignment.get(consumer1); List r3partitions2 = assignment.get(consumer2); List r3partitions3 = assignment.get(consumer3); @@ -811,7 +811,7 @@ public void testAssignmentWithConflictingPreviousGenerations() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(c3partitions0, 2)))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); List c3partitions = assignment.get(consumer3); @@ -851,7 +851,7 @@ public void testSchemaBackwardCompatibility() { new ConsumerUserData(c1partitions0, 1)))); subscriptions.put(consumer2, new Subscription(topics(topic), serializeTopicPartitionAssignmentToOldSchema(c2partitions0))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); List c3partitions = assignment.get(consumer3); @@ -890,13 +890,11 @@ public void testConflictingPreviousAssignments() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(c2partitions0, 1)))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); assertTrue(c1partitions.size() == 1 && c2partitions.size() == 1); - System.out.println(c1partitions); - System.out.println(c2partitions); assertEquals(2, assignor.generation()); verifyValidityAndBalance(subscriptions, assignment); assertTrue(isFullyBalanced(assignment)); From 4467d37e703a5dfeb2968761b6818a3881a0c2bb Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Wed, 25 Jul 2018 15:52:16 -0700 Subject: [PATCH 08/17] Minor updates --- .../clients/consumer/StickyAssignor.java | 32 ++++++++----------- .../internals/AbstractCoordinator.java | 2 +- .../internals/AbstractPartitionAssignor.java | 4 ++- .../clients/consumer/StickyAssignorTest.java | 20 ++++++------ 4 files changed, 27 insertions(+), 31 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index a36eb17c96a7d..c15d0c5ab33ca 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -30,7 +30,6 @@ import java.io.Serializable; import java.nio.ByteBuffer; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.Comparator; @@ -241,12 +240,12 @@ public Map> assign(Map partitionsP // 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()); + partition2AllPotentialConsumers.put(new TopicPartition(entry.getKey(), i), new ArrayList<>()); } for (Entry entry: subscriptions.entrySet()) { String consumer = entry.getKey(); - consumer2AllPotentialPartitions.put(consumer, new ArrayList()); + consumer2AllPotentialPartitions.put(consumer, new ArrayList<>()); 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); @@ -257,7 +256,7 @@ public Map> assign(Map partitionsP // add this consumer to currentAssignment (with an empty topic partition assignment) if it does not already exist if (!currentAssignment.containsKey(consumer)) - currentAssignment.put(consumer, new ArrayList()); + currentAssignment.put(consumer, new ArrayList<>()); } // a mapping of partition to current consumer @@ -374,7 +373,7 @@ private void prepopulateCurrentAssignments(Map subscriptio if (currentAssignment.containsKey(consumer)) currentAssignment.get(consumer).add(partition); else { - List partitions = new ArrayList<>(Arrays.asList(partition)); + List partitions = new ArrayList<>(Collections.singletonList(partition)); currentAssignment.put(consumer, partitions); } } @@ -519,7 +518,7 @@ private List sortPartitions(Map> cu TreeSet sortedConsumers = new TreeSet<>(new SubscriptionComparator(assignments)); sortedConsumers.addAll(assignments.keySet()); - List prevAssignedPartitions = new ArrayList(prevAssignment.keySet()); + List prevAssignedPartitions = new ArrayList<>(prevAssignment.keySet()); while (!sortedConsumers.isEmpty()) { String consumer = sortedConsumers.pollLast(); List remainingPartitions = assignments.get(consumer); @@ -564,17 +563,13 @@ private boolean areSubscriptionsIdentical(Map> part if (!hasIdenticalListElements(partition2AllPotentialConsumers.values())) return false; - if (!hasIdenticalListElements(consumer2AllPotentialPartitions.values())) - return false; - - return true; + return hasIdenticalListElements(consumer2AllPotentialPartitions.values()); } /** - * @return the consumer to which the given partition is assigned. The assignment should improve the overall balance - * of the partition assignments to consumers. + * The assignment should improve the overall balance of the partition assignments to consumers. */ - private String assignPartition(TopicPartition partition, + private void assignPartition(TopicPartition partition, TreeSet sortedCurrentSubscriptions, Map> currentAssignment, Map> consumer2AllPotentialPartitions, @@ -585,10 +580,9 @@ private String assignPartition(TopicPartition partition, currentAssignment.get(consumer).add(partition); currentPartitionConsumer.put(partition, consumer); sortedCurrentSubscriptions.add(consumer); - return consumer; + break; } } - return null; } private boolean canParticipateInReassignment(TopicPartition partition, @@ -632,7 +626,7 @@ private void balance(Map> currentAssignment, Map> partition2AllPotentialConsumers, Map currentPartitionConsumer) { boolean initializing = currentAssignment.get(sortedCurrentSubscriptions.last()).isEmpty(); - boolean reassignmentPerformed = false; + boolean reassignmentPerformed; // starts with a 'false' default value // assign all unassigned partitions for (TopicPartition partition: unassignedPartitions) { @@ -817,7 +811,7 @@ private static ConsumerUserData deserializeTopicPartitionAssignment(ByteBuffer b struct = STICKY_ASSIGNOR_USER_DATA_V0.read(copy); } catch (Exception e2) { // ignore the consumer's previous assignment if it cannot be parsed - return new ConsumerUserData(Arrays.asList(), -2); + return new ConsumerUserData(Collections.emptyList(), -2); } } @@ -930,11 +924,11 @@ private void addPartitionMovementRecord(TopicPartition partition, ConsumerPair p String topic = partition.topic(); if (!partitionMovementsByTopic.containsKey(topic)) - partitionMovementsByTopic.put(topic, new HashMap>()); + partitionMovementsByTopic.put(topic, new HashMap<>()); Map> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic); if (!partitionMovementsForThisTopic.containsKey(pair)) - partitionMovementsForThisTopic.put(pair, new HashSet()); + partitionMovementsForThisTopic.put(pair, new HashSet<>()); partitionMovementsForThisTopic.get(pair).add(partition); } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java index 4f29614943c22..8ddc837c47f90 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java @@ -583,7 +583,7 @@ private RequestFuture onJoinFollower() { // send follower's sync group with an empty assignment SyncGroupRequest.Builder requestBuilder = new SyncGroupRequest.Builder(groupId, generation.generationId, generation.memberId, - Collections.emptyMap()); + Collections.emptyMap()); log.debug("Sending follower SyncGroup to coordinator {}: {}", this.coordinator, requestBuilder); return sendSyncGroupRequest(requestBuilder); } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java index f7a06903f7761..d2a08b124bee1 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java @@ -59,7 +59,9 @@ public Map> assign(Map partitionsP Map subscriptions, int generation) { return assign(partitionsPerTopic, subscriptions); - } @Override + } + + @Override public Subscription subscription(Set topics) { return new Subscription(new ArrayList<>(topics)); } 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 e4253e5854158..a0d4edc355a0e 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 @@ -49,7 +49,7 @@ public void testOneConsumerNoTopic() { Map partitionsPerTopic = new HashMap<>(); Map subscriptions = - Collections.singletonMap(consumerId, new Subscription(Collections.emptyList())); + Collections.singletonMap(consumerId, new Subscription(Collections.emptyList())); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); assertEquals(Collections.singleton(consumerId), assignment.keySet()); @@ -369,7 +369,7 @@ public void testReassignmentAfterOneConsumerLeaves() { Map subscriptions = new HashMap<>(); for (int i = 1; i < 20; i++) { - List topics = new ArrayList(); + List topics = new ArrayList<>(); for (int j = 1; j <= i; j++) topics.add(getTopicName(j, 20)); subscriptions.put(getConsumerName(i, 20), new Subscription(topics)); @@ -419,7 +419,7 @@ public void testSameSubscriptions() { Map subscriptions = new HashMap<>(); for (int i = 1; i < 9; i++) { - List topics = new ArrayList(); + List topics = new ArrayList<>(); for (int j = 1; j <= partitionsPerTopic.size(); j++) topics.add(getTopicName(j, 15)); subscriptions.put(getConsumerName(i, 9), new Subscription(topics)); @@ -453,7 +453,7 @@ public void testLargeAssignmentWithMultipleConsumersLeaving() { Map subscriptions = new HashMap<>(); for (int i = 0; i < consumerCount; i++) { - List topics = new ArrayList(); + List topics = new ArrayList<>(); for (int j = 0; j < rand.nextInt(20); j++) topics.add(getTopicName(rand.nextInt(topicCount), topicCount)); subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics)); @@ -486,7 +486,7 @@ public void testNewSubscription() { Map subscriptions = new HashMap<>(); for (int i = 0; i < 3; i++) { - List topics = new ArrayList(); + List topics = new ArrayList<>(); for (int j = i; j <= 3 * i - 2; j++) topics.add(getTopicName(j, 5)); subscriptions.put(getConsumerName(i, 3), new Subscription(topics)); @@ -748,7 +748,7 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.remove(consumer3); assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); List r2partitions2 = assignment.get(consumer2); - assertTrue(r2partitions2.size() == 6); + assertEquals(6, r2partitions2.size()); assertEquals(2, assignor.generation()); assertTrue(r2partitions2.containsAll(r1partitions2)); verifyValidityAndBalance(subscriptions, assignment); @@ -770,9 +770,9 @@ public void testAssignmentWithMultipleGenerations2() { List r3partitions3 = assignment.get(consumer3); assertTrue(r3partitions1.size() == 2 && r3partitions2.size() == 2 && r3partitions3.size() == 2); assertEquals(3, assignor.generation()); - assertTrue(r3partitions1.equals(r1partitions1)); - assertTrue(r3partitions2.equals(r1partitions2)); - assertTrue(r3partitions3.equals(r1partitions3)); + assertEquals(r1partitions1, r3partitions1); + assertEquals(r1partitions2, r3partitions2); + assertEquals(r1partitions3, r3partitions3); verifyValidityAndBalance(subscriptions, assignment); assertTrue(isFullyBalanced(assignment)); assertTrue(assignor.isSticky()); @@ -901,7 +901,7 @@ public void testConflictingPreviousAssignments() { assertTrue(assignor.isSticky()); } - static ByteBuffer serializeTopicPartitionAssignmentToOldSchema(List partitions) { + private static ByteBuffer serializeTopicPartitionAssignmentToOldSchema(List partitions) { Struct struct = new Struct(StickyAssignor.STICKY_ASSIGNOR_USER_DATA_V0); List topicAssignments = new ArrayList<>(); for (Map.Entry> topicEntry : CollectionUtils.groupDataByTopic(partitions).entrySet()) { From 32eff14072f7535ac15de49da99aa7b6bd6e5a26 Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Wed, 1 Aug 2018 15:11:40 -0700 Subject: [PATCH 09/17] Fixed a couple of user data related issues --- .../kafka/clients/consumer/StickyAssignor.java | 14 ++++++++++++-- .../internals/AbstractPartitionAssignor.java | 4 ++++ .../kafka/clients/consumer/StickyAssignorTest.java | 2 +- 3 files changed, 17 insertions(+), 3 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index c15d0c5ab33ca..177b0da870607 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -353,13 +353,13 @@ private void prepopulateCurrentAssignments(Map subscriptio "following sticky assignment generation " + consumerUserData.generation + "."); } else if (consumerUserData.generation > existingRecord.generation) { // same partition is assigned to two consumers in different rebalances. - // the assignment being processed has higher generation that the one already processed, and takes priority + // the assignment being processed has higher generation than the one already processed, and takes priority updatePrevAssignment(prevAssignment, partition, existingRecord); partitionConsumer.put(partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); } else { // else if (consumerUserData.generation < existingRecord.generation) // if the same partition is assigned to two consumers in different rebalances and - // the assignment being processed has lower generation that the one already processed, + // the assignment being processed has lower generation than the one already processed, // it is only considered for updating the previous assignment of the involved topic partition updatePrevAssignment(prevAssignment, partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); } @@ -384,6 +384,7 @@ private void prepopulateCurrentAssignments(Map subscriptio @Override public void onAssignment(Assignment assignment) { memberAssignment = assignment.partitions(); + generation = deserializeTopicPartitionAssignment(assignment.userData()).generation; } @Override @@ -400,6 +401,15 @@ public String name() { return "sticky"; } + @Override + protected Map getAssignments(Map> rawAssignments) { + Map assignments = new HashMap<>(); + for (Map.Entry> assignmentEntry : rawAssignments.entrySet()) + assignments.put(assignmentEntry.getKey(), new Assignment(assignmentEntry.getValue(), + serializeTopicPartitionAssignment(new ConsumerUserData(assignmentEntry.getValue(), generation)))); + return assignments; + } + public int generation() { return generation; } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java index d2a08b124bee1..b9f1b8793c585 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java @@ -83,6 +83,10 @@ public Map assign(Cluster metadata, Map> rawAssignments = assign(partitionsPerTopic, subscriptions, generation); + return getAssignments(rawAssignments); + } + + protected Map getAssignments(Map> rawAssignments) { // this class maintains no user data, so just wrap the results Map assignments = new HashMap<>(); for (Map.Entry> assignmentEntry : rawAssignments.entrySet()) 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 a0d4edc355a0e..5f8ec4d518920 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 @@ -904,7 +904,7 @@ public void testConflictingPreviousAssignments() { private static ByteBuffer serializeTopicPartitionAssignmentToOldSchema(List partitions) { Struct struct = new Struct(StickyAssignor.STICKY_ASSIGNOR_USER_DATA_V0); List topicAssignments = new ArrayList<>(); - for (Map.Entry> topicEntry : CollectionUtils.groupDataByTopic(partitions).entrySet()) { + for (Map.Entry> topicEntry : CollectionUtils.groupPartitionsByTopic(partitions).entrySet()) { Struct topicAssignment = new Struct(StickyAssignor.TOPIC_ASSIGNMENT); topicAssignment.set(StickyAssignor.TOPIC_KEY_NAME, topicEntry.getKey()); topicAssignment.set(StickyAssignor.PARTITIONS_KEY_NAME, topicEntry.getValue().toArray()); From 10564ab9fc043c9d75a1c48df02d6d25d47bde26 Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Wed, 2 Jan 2019 21:18:16 -0800 Subject: [PATCH 10/17] Minor updates based on feedbacky --- .../apache/kafka/clients/consumer/StickyAssignor.java | 10 +++++----- .../consumer/internals/AbstractPartitionAssignor.java | 5 +++++ .../clients/consumer/internals/PartitionAssignor.java | 4 +--- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index 177b0da870607..f98433fb26c4a 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -203,8 +203,8 @@ public class StickyAssignor extends AbstractPartitionAssignor { private int generation; // consumer group generation static final class ConsumerUserData { - List partitions; - int generation; + final List partitions; + final int generation; ConsumerUserData(List partitions, int generation) { this.partitions = partitions; this.generation = generation; @@ -212,8 +212,8 @@ static final class ConsumerUserData { } static final class ConsumerGenerationPair { - String consumer; - int generation; + final String consumer; + final int generation; ConsumerGenerationPair(String consumer, int generation) { this.consumer = consumer; this.generation = generation; @@ -410,7 +410,7 @@ protected Map getAssignments(Map topics) { return new Subscription(new ArrayList<>(topics)); } + @Override + public Map assign(Cluster metadata, Map subscriptions) { + return assign(metadata, subscriptions, 0); + } + @Override public Map assign(Cluster metadata, Map subscriptions, int generation) { Set allSubscribedTopics = new HashSet<>(); diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java index 93d43a104205d..99a86ff089ee1 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java @@ -56,9 +56,7 @@ public interface PartitionAssignor { * @return A map from the members to their respective assignment. This should have one entry * for all members who in the input subscription map. */ - default Map assign(Cluster metadata, Map subscriptions) { - return assign(metadata, subscriptions, 0); - } + Map assign(Cluster metadata, Map subscriptions); /** * Perform the group assignment given the member subscriptions, current cluster metadata and group generation. From aca44c71dac29d8ccfdbc373d7df6f7349d9bf7c Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Mon, 7 Jan 2019 21:38:38 -0800 Subject: [PATCH 11/17] Removed the generation bump in StickyAssignor --- .../kafka/clients/consumer/StickyAssignor.java | 16 ++++++++-------- .../clients/consumer/StickyAssignorTest.java | 18 +++++++++--------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index f98433fb26c4a..982f4ae2046d1 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -334,7 +334,9 @@ private void prepopulateCurrentAssignments(Map subscriptio // higher generations overwrite lower generations in case of a conflict // note that a conflict could exists only if user data is for different generations - Map partitionConsumer = new HashMap<>(); + // currentPartitionConsumer will hold the latest ConsumerGenerationPair associated with a partition + // while prevAssignment holds the prior ConsumerGenerationPair (before current) of each parition + Map currentPartitionConsumer = new HashMap<>(); for (Map.Entry subscriptionEntry : subscriptions.entrySet()) { String consumer = subscriptionEntry.getKey(); ByteBuffer userData = subscriptionEntry.getValue().userData(); @@ -342,10 +344,10 @@ private void prepopulateCurrentAssignments(Map subscriptio ConsumerUserData consumerUserData = deserializeTopicPartitionAssignment(userData); for (TopicPartition partition: consumerUserData.partitions) { - if (!partitionConsumer.containsKey(partition)) - partitionConsumer.put(partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); + if (!currentPartitionConsumer.containsKey(partition)) + currentPartitionConsumer.put(partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); else { - ConsumerGenerationPair existingRecord = partitionConsumer.get(partition); + ConsumerGenerationPair existingRecord = currentPartitionConsumer.get(partition); if (consumerUserData.generation == existingRecord.generation) { // same partition is assigned to two consumers during the same rebalance. // log an error and skip this record @@ -355,7 +357,7 @@ private void prepopulateCurrentAssignments(Map subscriptio // same partition is assigned to two consumers in different rebalances. // the assignment being processed has higher generation than the one already processed, and takes priority updatePrevAssignment(prevAssignment, partition, existingRecord); - partitionConsumer.put(partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); + currentPartitionConsumer.put(partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); } else { // else if (consumerUserData.generation < existingRecord.generation) // if the same partition is assigned to two consumers in different rebalances and @@ -367,7 +369,7 @@ private void prepopulateCurrentAssignments(Map subscriptio } } - for (Map.Entry partitionAssignment: partitionConsumer.entrySet()) { + for (Map.Entry partitionAssignment: currentPartitionConsumer.entrySet()) { TopicPartition partition = partitionAssignment.getKey(); String consumer = partitionAssignment.getValue().consumer; if (currentAssignment.containsKey(consumer)) @@ -377,8 +379,6 @@ private void prepopulateCurrentAssignments(Map subscriptio currentAssignment.put(consumer, partitions); } } - - ++generation; } @Override 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 5f8ec4d518920..5161f5020f619 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 @@ -672,7 +672,7 @@ public void testAssignmentWithMultipleGenerations1() { subscriptions.put(consumer2, new Subscription(topics(topic))); subscriptions.put(consumer3, new Subscription(topics(topic))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 0); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); List r1partitions1 = assignment.get(consumer1); List r1partitions2 = assignment.get(consumer2); List r1partitions3 = assignment.get(consumer3); @@ -688,7 +688,7 @@ public void testAssignmentWithMultipleGenerations1() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions2, 1)))); subscriptions.remove(consumer3); - assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); + assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); List r2partitions1 = assignment.get(consumer1); List r2partitions2 = assignment.get(consumer2); assertTrue(r2partitions1.size() == 3 && r2partitions2.size() == 3); @@ -707,7 +707,7 @@ public void testAssignmentWithMultipleGenerations1() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions3, 1)))); - assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); + assignment = assignor.assign(partitionsPerTopic, subscriptions, 3); List r3partitions2 = assignment.get(consumer2); List r3partitions3 = assignment.get(consumer3); assertTrue(r3partitions2.size() == 3 && r3partitions3.size() == 3); @@ -732,7 +732,7 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.put(consumer2, new Subscription(topics(topic))); subscriptions.put(consumer3, new Subscription(topics(topic))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 0); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); List r1partitions1 = assignment.get(consumer1); List r1partitions2 = assignment.get(consumer2); List r1partitions3 = assignment.get(consumer3); @@ -746,7 +746,7 @@ public void testAssignmentWithMultipleGenerations2() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions2, 1)))); subscriptions.remove(consumer3); - assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); + assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); List r2partitions2 = assignment.get(consumer2); assertEquals(6, r2partitions2.size()); assertEquals(2, assignor.generation()); @@ -764,7 +764,7 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions3, 1)))); - assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); + assignment = assignor.assign(partitionsPerTopic, subscriptions, 3); List r3partitions1 = assignment.get(consumer1); List r3partitions2 = assignment.get(consumer2); List r3partitions3 = assignment.get(consumer3); @@ -811,7 +811,7 @@ public void testAssignmentWithConflictingPreviousGenerations() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(c3partitions0, 2)))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 3); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); List c3partitions = assignment.get(consumer3); @@ -851,7 +851,7 @@ public void testSchemaBackwardCompatibility() { new ConsumerUserData(c1partitions0, 1)))); subscriptions.put(consumer2, new Subscription(topics(topic), serializeTopicPartitionAssignmentToOldSchema(c2partitions0))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); List c3partitions = assignment.get(consumer3); @@ -890,7 +890,7 @@ public void testConflictingPreviousAssignments() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(c2partitions0, 1)))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); From 6ad099332e638e6dec2cf8e2885af671380d5049 Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Mon, 7 Jan 2019 23:40:06 -0800 Subject: [PATCH 12/17] Simplify finding current and last consumer of each partition by using a sorted map --- .../clients/consumer/StickyAssignor.java | 76 ++++++++----------- 1 file changed, 32 insertions(+), 44 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index 982f4ae2046d1..27db17d02440b 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -40,6 +40,7 @@ import java.util.Map; import java.util.Map.Entry; import java.util.Set; +import java.util.TreeMap; import java.util.TreeSet; /** @@ -310,23 +311,6 @@ public Map> assign(Map partitionsP return currentAssignment; } - private void updatePrevAssignment(Map prevAssignment, - TopicPartition partition, ConsumerGenerationPair consumerGeneration) { - if (!prevAssignment.containsKey(partition)) - prevAssignment.put(partition, consumerGeneration); - else { - ConsumerGenerationPair existingPair = prevAssignment.get(partition); - if (consumerGeneration.generation > existingPair.generation) - prevAssignment.put(partition, consumerGeneration); - else if (consumerGeneration.generation == existingPair.generation) { - // same partition is assigned to two consumers during the same rebalance. - // log an error and skip this record - log.error("Partition '" + partition + "' had been assigned to multiple consumers " + - "following sticky assignment generation " + consumerGeneration.generation + "."); - } - } - } - private void prepopulateCurrentAssignments(Map subscriptions, Map> currentAssignment, Map prevAssignment) { @@ -334,50 +318,54 @@ private void prepopulateCurrentAssignments(Map subscriptio // higher generations overwrite lower generations in case of a conflict // note that a conflict could exists only if user data is for different generations - // currentPartitionConsumer will hold the latest ConsumerGenerationPair associated with a partition - // while prevAssignment holds the prior ConsumerGenerationPair (before current) of each parition - Map currentPartitionConsumer = new HashMap<>(); - for (Map.Entry subscriptionEntry : subscriptions.entrySet()) { + // 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(); ByteBuffer userData = subscriptionEntry.getValue().userData(); if (userData == null || !userData.hasRemaining()) continue; ConsumerUserData consumerUserData = deserializeTopicPartitionAssignment(userData); for (TopicPartition partition: consumerUserData.partitions) { - if (!currentPartitionConsumer.containsKey(partition)) - currentPartitionConsumer.put(partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); - else { - ConsumerGenerationPair existingRecord = currentPartitionConsumer.get(partition); - if (consumerUserData.generation == existingRecord.generation) { + if (sortedPartitionConsumersByGeneration.containsKey(partition)) { + Map consumers = sortedPartitionConsumersByGeneration.get(partition); + if (consumers.containsKey(consumerUserData.generation)) { // same partition is assigned to two consumers during the same rebalance. - // log an error and skip this record - log.error("Partition '" + partition + "' is assigned to multiple consumers " + - "following sticky assignment generation " + consumerUserData.generation + "."); - } else if (consumerUserData.generation > existingRecord.generation) { - // same partition is assigned to two consumers in different rebalances. - // the assignment being processed has higher generation than the one already processed, and takes priority - updatePrevAssignment(prevAssignment, partition, existingRecord); - currentPartitionConsumer.put(partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); - } else { - // else if (consumerUserData.generation < existingRecord.generation) - // if the same partition is assigned to two consumers in different rebalances and - // the assignment being processed has lower generation than the one already processed, - // it is only considered for updating the previous assignment of the involved topic partition - updatePrevAssignment(prevAssignment, partition, new ConsumerGenerationPair(consumer, consumerUserData.generation)); - } + // log a warning and skip this record + log.warn("Partition '{}' is assigned to multiple consumers following sticky assignment generation {}.", + partition, consumerUserData.generation); + } else + consumers.put(consumerUserData.generation, consumer); + } else { + Map sortedConsumers = new TreeMap<>(); + sortedConsumers.put(consumerUserData.generation, consumer); + sortedPartitionConsumersByGeneration.put(partition, sortedConsumers); } } } - for (Map.Entry partitionAssignment: currentPartitionConsumer.entrySet()) { - TopicPartition partition = partitionAssignment.getKey(); - String consumer = partitionAssignment.getValue().consumer; + // 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 = (TreeMap) partitionConsumersEntry.getValue(); + Iterator it = consumers.descendingKeySet().iterator(); + + // let's process the current (most recent) consumer first + String consumer = consumers.get(it.next()); if (currentAssignment.containsKey(consumer)) currentAssignment.get(consumer).add(partition); else { List partitions = new ArrayList<>(Collections.singletonList(partition)); currentAssignment.put(consumer, partitions); } + + // now update previous assignment if any + if (it.hasNext()) { + int generation = it.next(); + prevAssignment.put(partition, new ConsumerGenerationPair(consumers.get(generation), generation)); + } + it.remove(); } } From e06a02b00c7a993df47c91e262d76ddfcba22b7f Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Tue, 8 Jan 2019 02:11:50 -0800 Subject: [PATCH 13/17] Made improvements based on the feedback --- .../clients/consumer/StickyAssignor.java | 43 +++++++++---------- .../internals/AbstractCoordinator.java | 5 ++- .../internals/AbstractPartitionAssignor.java | 7 +-- .../internals/ConsumerCoordinator.java | 5 ++- .../clients/consumer/StickyAssignorTest.java | 18 ++++---- .../internals/AbstractCoordinatorTest.java | 3 +- .../distributed/WorkerCoordinator.java | 2 +- 7 files changed, 42 insertions(+), 41 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index 27db17d02440b..7b151c04ba4b3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -39,6 +39,7 @@ import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.Optional; import java.util.Set; import java.util.TreeMap; import java.util.TreeSet; @@ -190,6 +191,7 @@ public class StickyAssignor extends AbstractPartitionAssignor { static final String TOPIC_KEY_NAME = "topic"; static final String PARTITIONS_KEY_NAME = "partitions"; private static final String GENERATION_KEY_NAME = "generation"; + private static final int DEFAULT_GENERATION = -1; static final Schema TOPIC_ASSIGNMENT = new Schema( new Field(TOPIC_KEY_NAME, Type.STRING), new Field(PARTITIONS_KEY_NAME, new ArrayOf(Type.INT32))); @@ -224,8 +226,8 @@ static final class ConsumerGenerationPair { @Override public Map> assign(Map partitionsPerTopic, Map subscriptions, - int generation) { - this.generation = generation; + Optional generation) { + this.generation = generation.orElse(DEFAULT_GENERATION); Map> currentAssignment = new HashMap<>(); Map prevAssignment = new HashMap<>(); partitionMovements = new PartitionMovements(); @@ -319,7 +321,7 @@ private void prepopulateCurrentAssignments(Map subscriptio // 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<>(); + Map> sortedPartitionConsumersByGeneration = new HashMap<>(); for (Map.Entry subscriptionEntry: subscriptions.entrySet()) { String consumer = subscriptionEntry.getKey(); ByteBuffer userData = subscriptionEntry.getValue().userData(); @@ -337,7 +339,7 @@ private void prepopulateCurrentAssignments(Map subscriptio } else consumers.put(consumerUserData.generation, consumer); } else { - Map sortedConsumers = new TreeMap<>(); + TreeMap sortedConsumers = new TreeMap<>(); sortedConsumers.put(consumerUserData.generation, consumer); sortedPartitionConsumersByGeneration.put(partition, sortedConsumers); } @@ -346,26 +348,21 @@ private void prepopulateCurrentAssignments(Map subscriptio // 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()) { + for (Map.Entry> partitionConsumersEntry: sortedPartitionConsumersByGeneration.entrySet()) { TopicPartition partition = partitionConsumersEntry.getKey(); - TreeMap consumers = (TreeMap) partitionConsumersEntry.getValue(); + TreeMap consumers = partitionConsumersEntry.getValue(); Iterator it = consumers.descendingKeySet().iterator(); // let's process the current (most recent) consumer first String consumer = consumers.get(it.next()); - if (currentAssignment.containsKey(consumer)) - currentAssignment.get(consumer).add(partition); - else { - List partitions = new ArrayList<>(Collections.singletonList(partition)); - currentAssignment.put(consumer, partitions); - } + 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)); } - it.remove(); } } @@ -389,14 +386,14 @@ public String name() { return "sticky"; } - @Override - protected Map getAssignments(Map> rawAssignments) { - Map assignments = new HashMap<>(); - for (Map.Entry> assignmentEntry : rawAssignments.entrySet()) - assignments.put(assignmentEntry.getKey(), new Assignment(assignmentEntry.getValue(), - serializeTopicPartitionAssignment(new ConsumerUserData(assignmentEntry.getValue(), generation)))); - return assignments; - } +// @Override +// protected Map getAssignments(Map> rawAssignments) { +// Map assignments = new HashMap<>(); +// for (Map.Entry> assignmentEntry : rawAssignments.entrySet()) +// assignments.put(assignmentEntry.getKey(), new Assignment(assignmentEntry.getValue(), +// serializeTopicPartitionAssignment(new ConsumerUserData(assignmentEntry.getValue(), generation)))); +// return assignments; +// } int generation() { return generation; @@ -624,7 +621,7 @@ private void balance(Map> currentAssignment, Map> partition2AllPotentialConsumers, Map currentPartitionConsumer) { boolean initializing = currentAssignment.get(sortedCurrentSubscriptions.last()).isEmpty(); - boolean reassignmentPerformed; // starts with a 'false' default value + boolean reassignmentPerformed = false; // assign all unassigned partitions for (TopicPartition partition: unassignedPartitions) { @@ -823,7 +820,7 @@ private static ConsumerUserData deserializeTopicPartitionAssignment(ByteBuffer b } } // make sure this is backward compatible - int generation = struct.hasField(GENERATION_KEY_NAME) ? struct.getInt(GENERATION_KEY_NAME) : -1; + int generation = struct.hasField(GENERATION_KEY_NAME) ? struct.getInt(GENERATION_KEY_NAME) : DEFAULT_GENERATION; return new ConsumerUserData(partitions, generation); } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java index 8ddc837c47f90..1ce9c76ce29ae 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java @@ -25,7 +25,6 @@ import org.apache.kafka.common.errors.GroupMaxSizeReachedException; import org.apache.kafka.common.errors.IllegalGenerationException; import org.apache.kafka.common.errors.InterruptException; -import org.apache.kafka.common.errors.MemberIdRequiredException; import org.apache.kafka.common.errors.RebalanceInProgressException; import org.apache.kafka.common.errors.RetriableException; import org.apache.kafka.common.errors.UnknownMemberIdException; @@ -87,7 +86,7 @@ * * To leverage this protocol, an implementation must define the format of metadata provided by each * member for group registration in {@link #metadata()} and the format of the state assignment provided - * by the leader in {@link #performAssignment(String, String, List)} and becomes available to members in + * by the leader in {@link #performAssignment(String, String, List, int)} and becomes available to members in * {@link #onJoinComplete(int, String, String, ByteBuffer)}. * * Note on locking: this class shares state between the caller and a background thread which is @@ -199,11 +198,13 @@ public AbstractCoordinator(LogContext logContext, * of the group (e.g. to push partition assignments in the case of the new consumer) * @param leaderId The id of the leader (which is this member) * @param allMemberMetadata Metadata from all members of the group + * @param generation Generation of the group * @return A map from each member to their state assignment */ protected abstract Map performAssignment(String leaderId, String protocol, List allMemberMetadata); + int generation); /** * Invoked when a group member has successfully joined a group. If this call fails with an exception, diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java index dc0d1b00a1ed8..1e88ce4522613 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java @@ -26,6 +26,7 @@ import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; /** @@ -44,7 +45,7 @@ public abstract class AbstractPartitionAssignor implements PartitionAssignor { */ public Map> assign(Map partitionsPerTopic, Map subscriptions) { - return assign(partitionsPerTopic, subscriptions, 0); + return assign(partitionsPerTopic, subscriptions, Optional.of(0)); } /** @@ -57,7 +58,7 @@ public Map> assign(Map partitionsP */ public Map> assign(Map partitionsPerTopic, Map subscriptions, - int generation) { + Optional generation) { return assign(partitionsPerTopic, subscriptions); } @@ -86,7 +87,7 @@ public Map assign(Cluster metadata, Map> rawAssignments = assign(partitionsPerTopic, subscriptions, generation); + Map> rawAssignments = assign(partitionsPerTopic, subscriptions, Optional.ofNullable(generation)); return getAssignments(rawAssignments); } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java index b3964bbd10eaf..d5a5eaf892ff1 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java @@ -390,7 +390,8 @@ private void updateGroupSubscription(Set topics) { @Override protected Map performAssignment(String leaderId, String assignmentStrategy, - List allSubscriptions) { + List allSubscriptions, + int generation) { PartitionAssignor assignor = lookupAssignor(assignmentStrategy); if (assignor == null) throw new IllegalStateException("Coordinator selected invalid assignment protocol: " + assignmentStrategy); @@ -411,7 +412,7 @@ protected Map performAssignment(String leaderId, log.debug("Performing assignment using strategy {} with subscriptions {}", assignor.name(), subscriptions); - Map assignment = assignor.assign(metadata.fetch(), subscriptions, latestGeneration().generationId); + Map assignment = assignor.assign(metadata.fetch(), subscriptions, generation); // user-customized assignor may have created some topics that are not in the subscription list // and assign their partitions to the members; in this case we would like to update the leader's 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 5161f5020f619..c341bdef78797 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 @@ -672,7 +672,7 @@ public void testAssignmentWithMultipleGenerations1() { subscriptions.put(consumer2, new Subscription(topics(topic))); subscriptions.put(consumer3, new Subscription(topics(topic))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(1)); List r1partitions1 = assignment.get(consumer1); List r1partitions2 = assignment.get(consumer2); List r1partitions3 = assignment.get(consumer3); @@ -688,7 +688,7 @@ public void testAssignmentWithMultipleGenerations1() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions2, 1)))); subscriptions.remove(consumer3); - assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); + assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(2)); List r2partitions1 = assignment.get(consumer1); List r2partitions2 = assignment.get(consumer2); assertTrue(r2partitions1.size() == 3 && r2partitions2.size() == 3); @@ -707,7 +707,7 @@ public void testAssignmentWithMultipleGenerations1() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions3, 1)))); - assignment = assignor.assign(partitionsPerTopic, subscriptions, 3); + assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(3)); List r3partitions2 = assignment.get(consumer2); List r3partitions3 = assignment.get(consumer3); assertTrue(r3partitions2.size() == 3 && r3partitions3.size() == 3); @@ -732,7 +732,7 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.put(consumer2, new Subscription(topics(topic))); subscriptions.put(consumer3, new Subscription(topics(topic))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 1); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(1)); List r1partitions1 = assignment.get(consumer1); List r1partitions2 = assignment.get(consumer2); List r1partitions3 = assignment.get(consumer3); @@ -746,7 +746,7 @@ public void testAssignmentWithMultipleGenerations2() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions2, 1)))); subscriptions.remove(consumer3); - assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); + assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(2)); List r2partitions2 = assignment.get(consumer2); assertEquals(6, r2partitions2.size()); assertEquals(2, assignor.generation()); @@ -764,7 +764,7 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions3, 1)))); - assignment = assignor.assign(partitionsPerTopic, subscriptions, 3); + assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(3)); List r3partitions1 = assignment.get(consumer1); List r3partitions2 = assignment.get(consumer2); List r3partitions3 = assignment.get(consumer3); @@ -811,7 +811,7 @@ public void testAssignmentWithConflictingPreviousGenerations() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(c3partitions0, 2)))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 3); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(3)); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); List c3partitions = assignment.get(consumer3); @@ -851,7 +851,7 @@ public void testSchemaBackwardCompatibility() { new ConsumerUserData(c1partitions0, 1)))); subscriptions.put(consumer2, new Subscription(topics(topic), serializeTopicPartitionAssignmentToOldSchema(c2partitions0))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(2)); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); List c3partitions = assignment.get(consumer3); @@ -890,7 +890,7 @@ public void testConflictingPreviousAssignments() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(c2partitions0, 1)))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, 2); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(2)); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java index dc8ea4642d54a..8360b4688b226 100644 --- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java @@ -745,7 +745,8 @@ protected JoinGroupRequestData.JoinGroupRequestProtocolSet metadata() { @Override protected Map performAssignment(String leaderId, String protocol, - List allMemberMetadata) { + List allMemberMetadata, + int generation) { Map assignment = new HashMap<>(); for (JoinGroupResponseData.JoinGroupResponseMember member : allMemberMetadata) { assignment.put(member.memberId(), EMPTY_DATA); diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java index fa89a9def403e..91c57632b1b4c 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java @@ -169,7 +169,7 @@ protected void onJoinComplete(int generation, String memberId, String protocol, } @Override - protected Map performAssignment(String leaderId, String protocol, List allMemberMetadata) { + protected Map performAssignment(String leaderId, String protocol, List allMemberMetadata, int generation) { log.debug("Performing task assignment"); Map memberConfigs = new HashMap<>(); From 980d858e9254f7736f89cdddc52cfdf13952f7eb Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Tue, 8 Jan 2019 06:05:00 -0800 Subject: [PATCH 14/17] Pass on the generation from the coordinator only after assignment is complete --- .../clients/consumer/StickyAssignor.java | 22 ++++--------- .../internals/AbstractCoordinator.java | 10 ++---- .../internals/AbstractPartitionAssignor.java | 33 +++---------------- .../internals/ConsumerCoordinator.java | 8 ++--- .../consumer/internals/PartitionAssignor.java | 22 ++++--------- .../clients/consumer/StickyAssignorTest.java | 27 +++++---------- .../internals/AbstractCoordinatorTest.java | 3 +- .../distributed/WorkerCoordinator.java | 2 +- .../internals/StreamsPartitionAssignor.java | 3 +- .../StreamsPartitionAssignorTest.java | 5 +-- .../streams/tests/StreamsUpgradeTest.java | 5 +-- 11 files changed, 41 insertions(+), 99 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index 7b151c04ba4b3..5745edbda7804 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -203,7 +203,7 @@ public class StickyAssignor extends AbstractPartitionAssignor { private List memberAssignment = null; private PartitionMovements partitionMovements; - private int generation; // consumer group generation + private int generation = DEFAULT_GENERATION; // consumer group generation static final class ConsumerUserData { final List partitions; @@ -223,11 +223,8 @@ static final class ConsumerGenerationPair { } } - @Override public Map> assign(Map partitionsPerTopic, - Map subscriptions, - Optional generation) { - this.generation = generation.orElse(DEFAULT_GENERATION); + Map subscriptions) { Map> currentAssignment = new HashMap<>(); Map prevAssignment = new HashMap<>(); partitionMovements = new PartitionMovements(); @@ -367,9 +364,10 @@ private void prepopulateCurrentAssignments(Map subscriptio } @Override - public void onAssignment(Assignment assignment) { + public void onAssignment(Assignment assignment, Optional generation) { memberAssignment = assignment.partitions(); - generation = deserializeTopicPartitionAssignment(assignment.userData()).generation; + this.generation = generation.orElse(DEFAULT_GENERATION); + System.out.println("Got " + memberAssignment + " of gen " + this.generation); } @Override @@ -377,6 +375,7 @@ public Subscription subscription(Set topics) { if (memberAssignment == null) return new Subscription(new ArrayList<>(topics)); + System.out.println("Sending " + memberAssignment + " from gen " + this.generation); return new Subscription(new ArrayList<>(topics), serializeTopicPartitionAssignment(new ConsumerUserData(memberAssignment, generation))); } @@ -386,15 +385,6 @@ public String name() { return "sticky"; } -// @Override -// protected Map getAssignments(Map> rawAssignments) { -// Map assignments = new HashMap<>(); -// for (Map.Entry> assignmentEntry : rawAssignments.entrySet()) -// assignments.put(assignmentEntry.getKey(), new Assignment(assignmentEntry.getValue(), -// serializeTopicPartitionAssignment(new ConsumerUserData(assignmentEntry.getValue(), generation)))); -// return assignments; -// } - int generation() { return generation; } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java index 1ce9c76ce29ae..9226014ab9c26 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java @@ -86,7 +86,7 @@ * * To leverage this protocol, an implementation must define the format of metadata provided by each * member for group registration in {@link #metadata()} and the format of the state assignment provided - * by the leader in {@link #performAssignment(String, String, List, int)} and becomes available to members in + * by the leader in {@link #performAssignment(String, String, List)} and becomes available to members in * {@link #onJoinComplete(int, String, String, ByteBuffer)}. * * Note on locking: this class shares state between the caller and a background thread which is @@ -198,13 +198,11 @@ public AbstractCoordinator(LogContext logContext, * of the group (e.g. to push partition assignments in the case of the new consumer) * @param leaderId The id of the leader (which is this member) * @param allMemberMetadata Metadata from all members of the group - * @param generation Generation of the group * @return A map from each member to their state assignment */ protected abstract Map performAssignment(String leaderId, String protocol, List allMemberMetadata); - int generation); /** * Invoked when a group member has successfully joined a group. If this call fails with an exception, @@ -584,7 +582,7 @@ private RequestFuture onJoinFollower() { // send follower's sync group with an empty assignment SyncGroupRequest.Builder requestBuilder = new SyncGroupRequest.Builder(groupId, generation.generationId, generation.memberId, - Collections.emptyMap()); + Collections.emptyMap()); log.debug("Sending follower SyncGroup to coordinator {}: {}", this.coordinator, requestBuilder); return sendSyncGroupRequest(requestBuilder); } @@ -754,10 +752,6 @@ protected synchronized Generation generation() { return generation; } - protected Generation latestGeneration() { - return generation; - } - /** * Check whether given generation id is matching the record within current generation. * Only using in unit tests. diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java index 1e88ce4522613..f7e01fef680dc 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java @@ -43,24 +43,8 @@ public abstract class AbstractPartitionAssignor implements PartitionAssignor { * @param subscriptions Map from the memberId to their respective topic subscription * @return Map from each member to the list of partitions assigned to them. */ - public Map> assign(Map partitionsPerTopic, - Map subscriptions) { - return assign(partitionsPerTopic, subscriptions, Optional.of(0)); - } - - /** - * Perform the group assignment given the partition counts, member subscriptions, and consumer group generation - * @param partitionsPerTopic The number of partitions for each subscribed topic. Topics not in metadata will be excluded - * from this map. - * @param subscriptions Map from the memberId to their respective topic subscription - * @param generation Consumer group generation id - * @return Map from each member to the list of partitions assigned to them. - */ - public Map> assign(Map partitionsPerTopic, - Map subscriptions, - Optional generation) { - return assign(partitionsPerTopic, subscriptions); - } + public abstract Map> assign(Map partitionsPerTopic, + Map subscriptions); @Override public Subscription subscription(Set topics) { @@ -69,11 +53,6 @@ public Subscription subscription(Set topics) { @Override public Map assign(Cluster metadata, Map subscriptions) { - return assign(metadata, subscriptions, 0); - } - - @Override - public Map assign(Cluster metadata, Map subscriptions, int generation) { Set allSubscribedTopics = new HashSet<>(); for (Map.Entry subscriptionEntry : subscriptions.entrySet()) allSubscribedTopics.addAll(subscriptionEntry.getValue().topics()); @@ -87,12 +66,8 @@ public Map assign(Cluster metadata, Map> rawAssignments = assign(partitionsPerTopic, subscriptions, Optional.ofNullable(generation)); - - return getAssignments(rawAssignments); - } + Map> rawAssignments = assign(partitionsPerTopic, subscriptions); - protected Map getAssignments(Map> rawAssignments) { // this class maintains no user data, so just wrap the results Map assignments = new HashMap<>(); for (Map.Entry> assignmentEntry : rawAssignments.entrySet()) @@ -101,7 +76,7 @@ protected Map getAssignments(Map generation) { // this assignor maintains no internal state, so nothing to do } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java index d5a5eaf892ff1..c438a8ec22f22 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java @@ -60,6 +60,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.atomic.AtomicInteger; @@ -262,7 +263,7 @@ protected void onJoinComplete(int generation, maybeUpdateJoinedSubscription(assignedPartitions); // give the assignor a chance to update internal state based on the received assignment - assignor.onAssignment(assignment); + assignor.onAssignment(assignment, Optional.ofNullable(generation)); // reschedule the auto commit starting from now if (autoCommitEnabled) @@ -390,8 +391,7 @@ private void updateGroupSubscription(Set topics) { @Override protected Map performAssignment(String leaderId, String assignmentStrategy, - List allSubscriptions, - int generation) { + List allSubscriptions) { PartitionAssignor assignor = lookupAssignor(assignmentStrategy); if (assignor == null) throw new IllegalStateException("Coordinator selected invalid assignment protocol: " + assignmentStrategy); @@ -412,7 +412,7 @@ protected Map performAssignment(String leaderId, log.debug("Performing assignment using strategy {} with subscriptions {}", assignor.name(), subscriptions); - Map assignment = assignor.assign(metadata.fetch(), subscriptions, generation); + Map assignment = assignor.assign(metadata.fetch(), subscriptions); // user-customized assignor may have created some topics that are not in the subscription list // and assign their partitions to the members; in this case we would like to update the leader's diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java index 99a86ff089ee1..fadfff6535047 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java @@ -22,6 +22,7 @@ import java.nio.ByteBuffer; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; /** @@ -29,7 +30,7 @@ * {@link org.apache.kafka.clients.consumer.KafkaConsumer}. Members of the consumer group subscribe * to the topics they are interested in and forward their subscriptions to a Kafka broker serving * as the group coordinator. The coordinator selects one member to perform the group assignment and - * propagates the subscriptions of all members to it. Then {@link #assign(Cluster, Map, int)} is called + * propagates the subscriptions of all members to it. Then {@link #assign(Cluster, Map)} is called * to perform the assignment and the results are forwarded back to each respective members * * In some cases, it is useful to forward additional metadata to the assignor in order to make @@ -42,7 +43,7 @@ public interface PartitionAssignor { /** * Return a serializable object representing the local member's subscription. This can include * additional information as well (e.g. local host/rack information) which can be leveraged in - * {@link #assign(Cluster, Map, int)}. + * {@link #assign(Cluster, Map)}. * @param topics Topics subscribed to through {@link org.apache.kafka.clients.consumer.KafkaConsumer#subscribe(java.util.Collection)} * and variants * @return Non-null subscription with optional user data @@ -58,23 +59,12 @@ public interface PartitionAssignor { */ Map assign(Cluster metadata, Map subscriptions); - /** - * Perform the group assignment given the member subscriptions, current cluster metadata and group generation. - * @param metadata Current topic/broker metadata known by consumer - * @param subscriptions Subscriptions from all members provided through {@link #subscription(Set)} - * @param generation Current group generation id - * @return A map from the members to their respective assignment. This should have one entry - * for all members who in the input subscription map. - */ - default Map assign(Cluster metadata, Map subscriptions, int generation) { - return assign(metadata, subscriptions); - } - /** * Callback which is invoked when a group member receives its assignment from the leader. - * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map, int)} + * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map)} + * @param generation The consumer group generation associated with this partition assignment (optional) */ - void onAssignment(Assignment assignment); + void onAssignment(Assignment assignment, Optional generation); /** 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 c341bdef78797..59f729c3c3c53 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 @@ -672,12 +672,11 @@ public void testAssignmentWithMultipleGenerations1() { subscriptions.put(consumer2, new Subscription(topics(topic))); subscriptions.put(consumer3, new Subscription(topics(topic))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(1)); + 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); - assertEquals(1, assignor.generation()); verifyValidityAndBalance(subscriptions, assignment); assertTrue(isFullyBalanced(assignment)); @@ -688,11 +687,10 @@ public void testAssignmentWithMultipleGenerations1() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions2, 1)))); subscriptions.remove(consumer3); - assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(2)); + assignment = assignor.assign(partitionsPerTopic, subscriptions); List r2partitions1 = assignment.get(consumer1); List r2partitions2 = assignment.get(consumer2); assertTrue(r2partitions1.size() == 3 && r2partitions2.size() == 3); - assertEquals(2, assignor.generation()); assertTrue(r2partitions1.containsAll(r1partitions1)); assertTrue(r2partitions2.containsAll(r1partitions2)); verifyValidityAndBalance(subscriptions, assignment); @@ -707,11 +705,10 @@ public void testAssignmentWithMultipleGenerations1() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions3, 1)))); - assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(3)); + assignment = assignor.assign(partitionsPerTopic, subscriptions); List r3partitions2 = assignment.get(consumer2); List r3partitions3 = assignment.get(consumer3); assertTrue(r3partitions2.size() == 3 && r3partitions3.size() == 3); - assertEquals(3, assignor.generation()); assertTrue(Collections.disjoint(r3partitions2, r3partitions3)); verifyValidityAndBalance(subscriptions, assignment); assertTrue(isFullyBalanced(assignment)); @@ -732,12 +729,11 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.put(consumer2, new Subscription(topics(topic))); subscriptions.put(consumer3, new Subscription(topics(topic))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(1)); + 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); - assertEquals(1, assignor.generation()); verifyValidityAndBalance(subscriptions, assignment); assertTrue(isFullyBalanced(assignment)); @@ -746,10 +742,9 @@ public void testAssignmentWithMultipleGenerations2() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions2, 1)))); subscriptions.remove(consumer3); - assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(2)); + assignment = assignor.assign(partitionsPerTopic, subscriptions); List r2partitions2 = assignment.get(consumer2); assertEquals(6, r2partitions2.size()); - assertEquals(2, assignor.generation()); assertTrue(r2partitions2.containsAll(r1partitions2)); verifyValidityAndBalance(subscriptions, assignment); assertTrue(isFullyBalanced(assignment)); @@ -764,12 +759,11 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(r1partitions3, 1)))); - assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(3)); + 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(3, assignor.generation()); assertEquals(r1partitions1, r3partitions1); assertEquals(r1partitions2, r3partitions2); assertEquals(r1partitions3, r3partitions3); @@ -811,13 +805,12 @@ public void testAssignmentWithConflictingPreviousGenerations() { subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(c3partitions0, 2)))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(3)); + 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); - assertEquals(3, assignor.generation()); assertTrue(c1partitions0.containsAll(c1partitions)); assertTrue(c2partitions0.containsAll(c2partitions)); assertTrue(c3partitions0.containsAll(c3partitions)); @@ -851,13 +844,12 @@ public void testSchemaBackwardCompatibility() { new ConsumerUserData(c1partitions0, 1)))); subscriptions.put(consumer2, new Subscription(topics(topic), serializeTopicPartitionAssignmentToOldSchema(c2partitions0))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(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() == 1 && c2partitions.size() == 1 && c3partitions.size() == 1); - assertEquals(2, assignor.generation()); assertTrue(c1partitions0.containsAll(c1partitions)); assertTrue(c2partitions0.containsAll(c2partitions)); verifyValidityAndBalance(subscriptions, assignment); @@ -890,12 +882,11 @@ public void testConflictingPreviousAssignments() { new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( new ConsumerUserData(c2partitions0, 1)))); - Map> assignment = assignor.assign(partitionsPerTopic, subscriptions, java.util.Optional.of(2)); + Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); assertTrue(c1partitions.size() == 1 && c2partitions.size() == 1); - assertEquals(2, assignor.generation()); verifyValidityAndBalance(subscriptions, assignment); assertTrue(isFullyBalanced(assignment)); assertTrue(assignor.isSticky()); diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java index 8360b4688b226..dc8ea4642d54a 100644 --- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java @@ -745,8 +745,7 @@ protected JoinGroupRequestData.JoinGroupRequestProtocolSet metadata() { @Override protected Map performAssignment(String leaderId, String protocol, - List allMemberMetadata, - int generation) { + List allMemberMetadata) { Map assignment = new HashMap<>(); for (JoinGroupResponseData.JoinGroupResponseMember member : allMemberMetadata) { assignment.put(member.memberId(), EMPTY_DATA); diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java index 91c57632b1b4c..fa89a9def403e 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java @@ -169,7 +169,7 @@ protected void onJoinComplete(int generation, String memberId, String protocol, } @Override - protected Map performAssignment(String leaderId, String protocol, List allMemberMetadata, int generation) { + protected Map performAssignment(String leaderId, String protocol, List allMemberMetadata) { log.debug("Performing task assignment"); Map memberConfigs = new HashMap<>(); diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 3f3c1233da006..c5442c06f60b6 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -49,6 +49,7 @@ import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; @@ -797,7 +798,7 @@ List> interleaveTasksByGroupId(final Collection taskIds, fi * @throws TaskAssignmentException if there is no task id for one of the partitions specified */ @Override - public void onAssignment(final Assignment assignment) { + public void onAssignment(final Assignment assignment, final Optional generation) { final List partitions = new ArrayList<>(assignment.partitions()); Collections.sort(partitions, PARTITION_COMPARATOR); diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index 284e396f978c6..53f0fa0733482 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -53,6 +53,7 @@ import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; @@ -655,7 +656,7 @@ public void testOnAssignment() { EasyMock.expectLastCall(); EasyMock.replay(taskManager); - partitionAssignor.onAssignment(assignment); + partitionAssignor.onAssignment(assignment, Optional.empty()); EasyMock.verify(taskManager); @@ -976,7 +977,7 @@ public void shouldUpdateClusterMetadataAndHostInfoOnAssignment() { EasyMock.expectLastCall(); EasyMock.replay(taskManager); - partitionAssignor.onAssignment(createAssignment(hostState)); + partitionAssignor.onAssignment(createAssignment(hostState), Optional.empty()); EasyMock.verify(taskManager); } diff --git a/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java b/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java index 33e9b9771c245..1d8e958be6a13 100644 --- a/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java @@ -51,6 +51,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Properties; import java.util.Set; import java.util.UUID; @@ -141,9 +142,9 @@ public Subscription subscription(final Set topics) { } @Override - public void onAssignment(final PartitionAssignor.Assignment assignment) { + public void onAssignment(final PartitionAssignor.Assignment assignment, final Optional generation) { try { - super.onAssignment(assignment); + super.onAssignment(assignment, generation); return; } catch (final TaskAssignmentException cannotProcessFutureVersion) { // continue From f2be37d9c79b3ebd82fa4286f03802069978d982 Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Sat, 26 Jan 2019 00:13:38 -0800 Subject: [PATCH 15/17] Use `Optional` for `generation` to improve backward compatibility Also, minor fixes and clarifying comments --- .../clients/consumer/StickyAssignor.java | 42 ++++++++----- .../internals/AbstractPartitionAssignor.java | 7 ++- .../internals/ConsumerCoordinator.java | 2 +- .../consumer/internals/PartitionAssignor.java | 6 ++ .../clients/consumer/StickyAssignorTest.java | 61 ++++++++++--------- .../internals/StreamsPartitionAssignor.java | 7 ++- 6 files changed, 76 insertions(+), 49 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index 5745edbda7804..b07645caf2274 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -207,8 +207,8 @@ public class StickyAssignor extends AbstractPartitionAssignor { static final class ConsumerUserData { final List partitions; - final int generation; - ConsumerUserData(List partitions, int generation) { + final Optional generation; + ConsumerUserData(List partitions, Optional generation) { this.partitions = partitions; this.generation = generation; } @@ -266,7 +266,7 @@ public Map> assign(Map partitionsP currentPartitionConsumer.put(topicPartition, entry.getKey()); List sortedPartitions = sortPartitions( - currentAssignment, prevAssignment, isFreshAssignment, partition2AllPotentialConsumers, consumer2AllPotentialPartitions); + currentAssignment, prevAssignment.keySet(), isFreshAssignment, partition2AllPotentialConsumers, consumer2AllPotentialPartitions); // all partitions that need to be assigned (initially set to all partitions but adjusted in the following loop) List unassignedPartitions = new ArrayList<>(sortedPartitions); @@ -328,16 +328,16 @@ private void prepopulateCurrentAssignments(Map subscriptio for (TopicPartition partition: consumerUserData.partitions) { if (sortedPartitionConsumersByGeneration.containsKey(partition)) { Map consumers = sortedPartitionConsumersByGeneration.get(partition); - if (consumers.containsKey(consumerUserData.generation)) { + if (consumerUserData.generation.isPresent() && consumers.containsKey(consumerUserData.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, consumerUserData.generation); } else - consumers.put(consumerUserData.generation, consumer); + consumers.put(consumerUserData.generation.orElse(DEFAULT_GENERATION), consumer); } else { TreeMap sortedConsumers = new TreeMap<>(); - sortedConsumers.put(consumerUserData.generation, consumer); + sortedConsumers.put(consumerUserData.generation.orElse(DEFAULT_GENERATION), consumer); sortedPartitionConsumersByGeneration.put(partition, sortedConsumers); } } @@ -367,7 +367,6 @@ private void prepopulateCurrentAssignments(Map subscriptio public void onAssignment(Assignment assignment, Optional generation) { memberAssignment = assignment.partitions(); this.generation = generation.orElse(DEFAULT_GENERATION); - System.out.println("Got " + memberAssignment + " of gen " + this.generation); } @Override @@ -375,9 +374,8 @@ public Subscription subscription(Set topics) { if (memberAssignment == null) return new Subscription(new ArrayList<>(topics)); - System.out.println("Sending " + memberAssignment + " from gen " + this.generation); return new Subscription(new ArrayList<>(topics), - serializeTopicPartitionAssignment(new ConsumerUserData(memberAssignment, generation))); + serializeTopicPartitionAssignment(new ConsumerUserData(memberAssignment, Optional.of(generation)))); } @Override @@ -474,14 +472,16 @@ private int getBalanceScore(Map> assignment) { * that causes minimal partition movement among consumers (hence honoring maximal stickiness) * * @param currentAssignment the calculated assignment so far - * @param prevAssignment previous consumer of each partition, if any + * @param partitionsWithADifferentPreviousAssignment partitions that had a different consumer before (for every + * such partition there should also be a mapping in + * @currentAssignment to a different consumer) * @param isFreshAssignment whether this is a new assignment, or a reassignment of an existing one * @param partition2AllPotentialConsumers a mapping of partitions to their potential consumers * @param consumer2AllPotentialPartitions a mapping of consumers to potential partitions they can consumer from * @return sorted list of valid partitions */ private List sortPartitions(Map> currentAssignment, - Map prevAssignment, + Set partitionsWithADifferentPreviousAssignment, boolean isFreshAssignment, Map> partition2AllPotentialConsumers, Map> consumer2AllPotentialPartitions) { @@ -502,19 +502,28 @@ private List sortPartitions(Map> cu } TreeSet sortedConsumers = new TreeSet<>(new SubscriptionComparator(assignments)); sortedConsumers.addAll(assignments.keySet()); + // at this point, sortedConsumers contains an ascending-sorted list of consumers based on + // how many valid partitions are currently assigned to them - List prevAssignedPartitions = new ArrayList<>(prevAssignment.keySet()); while (!sortedConsumers.isEmpty()) { + // take the consumer with the most partitions String consumer = sortedConsumers.pollLast(); + // currently assigned partitions to this consumer List remainingPartitions = assignments.get(consumer); - List prevPartitions = new ArrayList<>(prevAssignedPartitions); + // partitions that were assigned to a different consumer last time + List prevPartitions = new ArrayList<>(partitionsWithADifferentPreviousAssignment); + // from partitions that had a different consumer before, keep only those that are + // assigned to this consumer now prevPartitions.retainAll(remainingPartitions); if (!prevPartitions.isEmpty()) { + // if there is a partition of this consumer that was assigned to another consumer before + // mark it as good options for reassignment TopicPartition partition = prevPartitions.remove(0); remainingPartitions.remove(partition); sortedPartitions.add(partition); sortedConsumers.add(consumer); } else if (!remainingPartitions.isEmpty()) { + // otherwise, mark any other one of the current partitions as a reassignment candidate sortedPartitions.add(remainingPartitions.remove(0)); sortedConsumers.add(consumer); } @@ -778,7 +787,8 @@ static ByteBuffer serializeTopicPartitionAssignment(ConsumerUserData consumerUse topicAssignments.add(topicAssignment); } struct.set(TOPIC_PARTITIONS_KEY_NAME, topicAssignments.toArray()); - struct.set(GENERATION_KEY_NAME, consumerUserData.generation); + if (consumerUserData.generation.isPresent()) + struct.set(GENERATION_KEY_NAME, consumerUserData.generation.get()); ByteBuffer buffer = ByteBuffer.allocate(STICKY_ASSIGNOR_USER_DATA_V1.sizeOf(struct)); STICKY_ASSIGNOR_USER_DATA_V1.write(buffer, struct); buffer.flip(); @@ -796,7 +806,7 @@ private static ConsumerUserData deserializeTopicPartitionAssignment(ByteBuffer b struct = STICKY_ASSIGNOR_USER_DATA_V0.read(copy); } catch (Exception e2) { // ignore the consumer's previous assignment if it cannot be parsed - return new ConsumerUserData(Collections.emptyList(), -2); + return new ConsumerUserData(Collections.emptyList(), Optional.of(DEFAULT_GENERATION)); } } @@ -810,7 +820,7 @@ private static ConsumerUserData deserializeTopicPartitionAssignment(ByteBuffer b } } // make sure this is backward compatible - int generation = struct.hasField(GENERATION_KEY_NAME) ? struct.getInt(GENERATION_KEY_NAME) : DEFAULT_GENERATION; + Optional generation = struct.hasField(GENERATION_KEY_NAME) ? Optional.of(struct.getInt(GENERATION_KEY_NAME)) : Optional.empty(); return new ConsumerUserData(partitions, generation); } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java index f7e01fef680dc..9e5dab6a6fb2f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java @@ -76,10 +76,15 @@ public Map assign(Cluster metadata, Map generation) { + public void onAssignment(Assignment assignment) { // this assignor maintains no internal state, so nothing to do } + @Override + public void onAssignment(Assignment assignment, Optional generation) { + onAssignment(assignment); + } + protected static void put(Map> map, K key, V value) { List list = map.computeIfAbsent(key, k -> new ArrayList<>()); list.add(value); diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java index c438a8ec22f22..7206144a098d3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java @@ -263,7 +263,7 @@ protected void onJoinComplete(int generation, maybeUpdateJoinedSubscription(assignedPartitions); // give the assignor a chance to update internal state based on the received assignment - assignor.onAssignment(assignment, Optional.ofNullable(generation)); + assignor.onAssignment(assignment, Optional.of(generation)); // reschedule the auto commit starting from now if (autoCommitEnabled) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java index fadfff6535047..59ab39b14deac 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java @@ -59,6 +59,12 @@ public interface PartitionAssignor { */ Map assign(Cluster metadata, Map subscriptions); + /** + * Callback which is invoked when a group member receives its assignment from the leader. + * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map)} + */ + void onAssignment(Assignment assignment); + /** * Callback which is invoked when a group member receives its assignment from the leader. * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map)} 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 59f729c3c3c53..3e67551eac1a1 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 @@ -28,6 +28,7 @@ import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.Optional; import java.util.Random; import java.util.Set; @@ -239,7 +240,7 @@ public void testAddRemoveConsumerOneTopic() { String consumer2 = "consumer2"; subscriptions.put(consumer1, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(assignment.get(consumer1), assignor.generation())))); + new ConsumerUserData(assignment.get(consumer1), Optional.of(assignor.generation()))))); subscriptions.put(consumer2, new Subscription(topics(topic))); assignment = assignor.assign(partitionsPerTopic, subscriptions); assertEquals(partitions(tp(topic, 2), tp(topic, 1)), assignment.get(consumer1)); @@ -252,7 +253,7 @@ public void testAddRemoveConsumerOneTopic() { subscriptions.remove(consumer1); subscriptions.put(consumer2, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(assignment.get(consumer2), assignor.generation())))); + new ConsumerUserData(assignment.get(consumer2), Optional.of(assignor.generation()))))); assignment = assignor.assign(partitionsPerTopic, subscriptions); assertTrue(assignment.get(consumer2).contains(tp(topic, 0))); assertTrue(assignment.get(consumer2).contains(tp(topic, 1))); @@ -324,10 +325,10 @@ public void testAddRemoveTopicTwoConsumers() { partitionsPerTopic.put(topic2, 3); subscriptions.put(consumer1, new Subscription(topics(topic, topic2), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(assignment.get(consumer1), assignor.generation())))); + new ConsumerUserData(assignment.get(consumer1), Optional.of(assignor.generation()))))); subscriptions.put(consumer2, new Subscription(topics(topic, topic2), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(assignment.get(consumer2), assignor.generation())))); + new ConsumerUserData(assignment.get(consumer2), Optional.of(assignor.generation()))))); assignment = assignor.assign(partitionsPerTopic, subscriptions); // verify balance verifyValidityAndBalance(subscriptions, assignment); @@ -343,10 +344,10 @@ public void testAddRemoveTopicTwoConsumers() { partitionsPerTopic.remove(topic); subscriptions.put(consumer1, new Subscription(topics(topic2), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(assignment.get(consumer1), assignor.generation())))); + new ConsumerUserData(assignment.get(consumer1), Optional.of(assignor.generation()))))); subscriptions.put(consumer2, new Subscription(topics(topic2), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(assignment.get(consumer2), assignor.generation())))); + new ConsumerUserData(assignment.get(consumer2), Optional.of(assignor.generation()))))); assignment = assignor.assign(partitionsPerTopic, subscriptions); // verify balance verifyValidityAndBalance(subscriptions, assignment); @@ -382,7 +383,7 @@ public void testReassignmentAfterOneConsumerLeaves() { String consumer = getConsumerName(i, 20); subscriptions.put(consumer, new Subscription(subscriptions.get(consumer).topics(), - StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), assignor.generation())))); + StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), Optional.of(assignor.generation()))))); } subscriptions.remove("consumer10"); @@ -432,7 +433,7 @@ public void testSameSubscriptions() { String consumer = getConsumerName(i, 9); subscriptions.put(consumer, new Subscription(subscriptions.get(consumer).topics(), - StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), assignor.generation())))); + StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), Optional.of(assignor.generation()))))); } subscriptions.remove(getConsumerName(5, 9)); @@ -466,7 +467,7 @@ public void testLargeAssignmentWithMultipleConsumersLeaving() { String consumer = getConsumerName(i, consumerCount); subscriptions.put(consumer, new Subscription(subscriptions.get(consumer).topics(), - StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), assignor.generation())))); + StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), Optional.of(assignor.generation()))))); } for (int i = 0; i < 50; ++i) { String c = getConsumerName(rand.nextInt(consumerCount), consumerCount); @@ -538,7 +539,7 @@ public void testReassignmentWithRandomSubscriptionsAndChanges() { List sub = Utils.sorted(getRandomSublist(topics)); String consumer = getConsumerName(i, maxNumConsumers); subscriptions.put(consumer, - new Subscription(sub, StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), assignor.generation())))); + new Subscription(sub, StickyAssignor.serializeTopicPartitionAssignment(new ConsumerUserData(assignment.get(consumer), Optional.of(assignor.generation()))))); } assignment = assignor.assign(partitionsPerTopic, subscriptions); @@ -557,15 +558,15 @@ public void testMoveExistingAssignments() { subscriptions.put("consumer01", new Subscription(topics("topic01", "topic02"), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(partitions(tp("topic01", 0)), assignor.generation())))); + new ConsumerUserData(partitions(tp("topic01", 0)), Optional.of(assignor.generation()))))); subscriptions.put("consumer02", new Subscription(topics("topic01", "topic02", "topic03", "topic04"), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(partitions(tp("topic02", 0), tp("topic03", 0)), assignor.generation())))); + new ConsumerUserData(partitions(tp("topic02", 0), tp("topic03", 0)), Optional.of(assignor.generation()))))); subscriptions.put("consumer03", new Subscription(topics("topic02", "topic03", "topic04", "topic05", "topic06"), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(partitions(tp("topic04", 0), tp("topic05", 0), tp("topic06", 0)), assignor.generation())))); + new ConsumerUserData(partitions(tp("topic04", 0), tp("topic05", 0), tp("topic06", 0)), Optional.of(assignor.generation()))))); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment); @@ -600,15 +601,15 @@ public void testStickiness() { subscriptions.put("consumer02", new Subscription(topics("topic01"), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(assignment.get("consumer02"), assignor.generation())))); + new ConsumerUserData(assignment.get("consumer02"), Optional.of(assignor.generation()))))); subscriptions.put("consumer03", new Subscription(topics("topic01"), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(assignment.get("consumer03"), assignor.generation())))); + new ConsumerUserData(assignment.get("consumer03"), Optional.of(assignor.generation()))))); subscriptions.put("consumer04", new Subscription(topics("topic01"), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(assignment.get("consumer04"), assignor.generation())))); + new ConsumerUserData(assignment.get("consumer04"), Optional.of(assignor.generation()))))); assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment); @@ -682,10 +683,10 @@ public void testAssignmentWithMultipleGenerations1() { subscriptions.put(consumer1, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(r1partitions1, 1)))); + new ConsumerUserData(r1partitions1, Optional.of(1))))); subscriptions.put(consumer2, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(r1partitions2, 1)))); + new ConsumerUserData(r1partitions2, Optional.of(1))))); subscriptions.remove(consumer3); assignment = assignor.assign(partitionsPerTopic, subscriptions); List r2partitions1 = assignment.get(consumer1); @@ -701,10 +702,10 @@ public void testAssignmentWithMultipleGenerations1() { subscriptions.remove(consumer1); subscriptions.put(consumer2, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(r2partitions2, 2)))); + new ConsumerUserData(r2partitions2, Optional.of(2))))); subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(r1partitions3, 1)))); + new ConsumerUserData(r1partitions3, Optional.of(1))))); assignment = assignor.assign(partitionsPerTopic, subscriptions); List r3partitions2 = assignment.get(consumer2); List r3partitions3 = assignment.get(consumer3); @@ -740,7 +741,7 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.remove(consumer1); subscriptions.put(consumer2, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(r1partitions2, 1)))); + new ConsumerUserData(r1partitions2, Optional.of(1))))); subscriptions.remove(consumer3); assignment = assignor.assign(partitionsPerTopic, subscriptions); List r2partitions2 = assignment.get(consumer2); @@ -752,13 +753,13 @@ public void testAssignmentWithMultipleGenerations2() { subscriptions.put(consumer1, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(r1partitions1, 1)))); + new ConsumerUserData(r1partitions1, Optional.of(1))))); subscriptions.put(consumer2, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(r2partitions2, 2)))); + new ConsumerUserData(r2partitions2, Optional.of(2))))); subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(r1partitions3, 1)))); + new ConsumerUserData(r1partitions3, Optional.of(1))))); assignment = assignor.assign(partitionsPerTopic, subscriptions); List r3partitions1 = assignment.get(consumer1); List r3partitions2 = assignment.get(consumer2); @@ -798,13 +799,13 @@ public void testAssignmentWithConflictingPreviousGenerations() { List c3partitions0 = partitions(tp3, tp4, tp5); subscriptions.put(consumer1, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(c1partitions0, 1)))); + new ConsumerUserData(c1partitions0, Optional.of(1))))); subscriptions.put(consumer2, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(c2partitions0, 1)))); + new ConsumerUserData(c2partitions0, Optional.of(1))))); subscriptions.put(consumer3, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(c3partitions0, 2)))); + new ConsumerUserData(c3partitions0, Optional.of(2))))); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); List c1partitions = assignment.get(consumer1); List c2partitions = assignment.get(consumer2); @@ -841,7 +842,7 @@ public void testSchemaBackwardCompatibility() { List c2partitions0 = partitions(tp1); subscriptions.put(consumer1, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(c1partitions0, 1)))); + new ConsumerUserData(c1partitions0, Optional.of(1))))); subscriptions.put(consumer2, new Subscription(topics(topic), serializeTopicPartitionAssignmentToOldSchema(c2partitions0))); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); @@ -877,10 +878,10 @@ public void testConflictingPreviousAssignments() { List c2partitions0 = partitions(tp0, tp1); subscriptions.put(consumer1, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(c1partitions0, 1)))); + new ConsumerUserData(c1partitions0, Optional.of(1))))); subscriptions.put(consumer2, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(c2partitions0, 1)))); + new ConsumerUserData(c2partitions0, Optional.of(1))))); Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); List c1partitions = assignment.get(consumer1); diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index c5442c06f60b6..bf5027f57fd5a 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -794,11 +794,16 @@ List> interleaveTasksByGroupId(final Collection taskIds, fi return taskIdsForConsumerAssignment; } + @Override + public void onAssignment(final Assignment assignment, final Optional generation) { + onAssignment(assignment); + } + /** * @throws TaskAssignmentException if there is no task id for one of the partitions specified */ @Override - public void onAssignment(final Assignment assignment, final Optional generation) { + public void onAssignment(final Assignment assignment) { final List partitions = new ArrayList<>(assignment.partitions()); Collections.sort(partitions, PARTITION_COMPARATOR); From 2e410c08ae9b28595d14a7ba3046deb20b1739c9 Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Sat, 23 Feb 2019 21:53:26 -0800 Subject: [PATCH 16/17] Added a default implementation to the interface --- .../consumer/internals/AbstractPartitionAssignor.java | 6 ------ .../kafka/clients/consumer/internals/PartitionAssignor.java | 4 +++- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java index 9e5dab6a6fb2f..35eb8eb141f93 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java @@ -26,7 +26,6 @@ import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Set; /** @@ -80,11 +79,6 @@ public void onAssignment(Assignment assignment) { // this assignor maintains no internal state, so nothing to do } - @Override - public void onAssignment(Assignment assignment, Optional generation) { - onAssignment(assignment); - } - protected static void put(Map> map, K key, V value) { List list = map.computeIfAbsent(key, k -> new ArrayList<>()); list.add(value); diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java index 59ab39b14deac..59c2e9b69ba3d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java @@ -70,7 +70,9 @@ public interface PartitionAssignor { * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map)} * @param generation The consumer group generation associated with this partition assignment (optional) */ - void onAssignment(Assignment assignment, Optional generation); + default void onAssignment(Assignment assignment, Optional generation) { + onAssignment(assignment); + } /** From a03a73432c4febeeb166b18d07fe24a304676a44 Mon Sep 17 00:00:00 2001 From: Vahid Hashemian Date: Sat, 13 Apr 2019 12:48:02 -0700 Subject: [PATCH 17/17] Removed the optional qualifier from `onAssignment`'s 'generation` argument --- .../kafka/clients/consumer/StickyAssignor.java | 4 ++-- .../consumer/internals/AbstractCoordinator.java | 13 +++++++------ .../consumer/internals/ConsumerCoordinator.java | 3 +-- .../consumer/internals/PartitionAssignor.java | 3 +-- .../kafka/clients/consumer/StickyAssignorTest.java | 2 +- .../internals/StreamsPartitionAssignor.java | 6 ------ .../internals/StreamsPartitionAssignorTest.java | 5 ++--- .../kafka/streams/tests/StreamsUpgradeTest.java | 5 ++--- 8 files changed, 16 insertions(+), 25 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index b07645caf2274..9575ba6dc42e0 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -364,9 +364,9 @@ private void prepopulateCurrentAssignments(Map subscriptio } @Override - public void onAssignment(Assignment assignment, Optional generation) { + public void onAssignment(Assignment assignment, int generation) { memberAssignment = assignment.partitions(); - this.generation = generation.orElse(DEFAULT_GENERATION); + this.generation = generation; } @Override diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java index 9226014ab9c26..9261966ccbd39 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java @@ -25,6 +25,7 @@ import org.apache.kafka.common.errors.GroupMaxSizeReachedException; import org.apache.kafka.common.errors.IllegalGenerationException; import org.apache.kafka.common.errors.InterruptException; +import org.apache.kafka.common.errors.MemberIdRequiredException; import org.apache.kafka.common.errors.RebalanceInProgressException; import org.apache.kafka.common.errors.RetriableException; import org.apache.kafka.common.errors.UnknownMemberIdException; @@ -565,7 +566,7 @@ public void handle(JoinGroupResponse joinResponse, RequestFuture fut // and send another join group request in next cycle. synchronized (AbstractCoordinator.this) { AbstractCoordinator.this.generation = new Generation(OffsetCommitRequest.DEFAULT_GENERATION_ID, - joinResponse.data().memberId(), null); + joinResponse.data().memberId(), null); AbstractCoordinator.this.rejoinNeeded = true; AbstractCoordinator.this.state = MemberState.UNJOINED; } @@ -653,7 +654,7 @@ private RequestFuture sendFindCoordinatorRequest(Node node) { FindCoordinatorRequest.Builder requestBuilder = new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, this.groupId); return client.send(node, requestBuilder) - .compose(new FindCoordinatorResponseHandler()); + .compose(new FindCoordinatorResponseHandler()); } private class FindCoordinatorResponseHandler extends RequestFutureAdapter { @@ -939,8 +940,8 @@ public GroupCoordinatorMetrics(Metrics metrics, String metricGrpPrefix) { this.heartbeatLatency = metrics.sensor("heartbeat-latency"); this.heartbeatLatency.add(metrics.metricName("heartbeat-response-time-max", - this.metricGrpName, - "The max time taken to receive a response to a heartbeat request"), new Max()); + this.metricGrpName, + "The max time taken to receive a response to a heartbeat request"), new Max()); this.heartbeatLatency.add(createMeter(metrics, metricGrpName, "heartbeat", "heartbeats")); this.joinLatency = metrics.sensor("join-latency"); @@ -1147,8 +1148,8 @@ public boolean equals(final Object o) { if (o == null || getClass() != o.getClass()) return false; final Generation that = (Generation) o; return generationId == that.generationId && - Objects.equals(memberId, that.memberId) && - Objects.equals(protocol, that.protocol); + Objects.equals(memberId, that.memberId) && + Objects.equals(protocol, that.protocol); } @Override diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java index 7206144a098d3..b31bf44ae7fb0 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java @@ -60,7 +60,6 @@ import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.Optional; import java.util.Set; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.atomic.AtomicInteger; @@ -263,7 +262,7 @@ protected void onJoinComplete(int generation, maybeUpdateJoinedSubscription(assignedPartitions); // give the assignor a chance to update internal state based on the received assignment - assignor.onAssignment(assignment, Optional.of(generation)); + assignor.onAssignment(assignment, generation); // reschedule the auto commit starting from now if (autoCommitEnabled) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java index 59c2e9b69ba3d..43fdaf3b01cd3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java @@ -22,7 +22,6 @@ import java.nio.ByteBuffer; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Set; /** @@ -70,7 +69,7 @@ public interface PartitionAssignor { * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map)} * @param generation The consumer group generation associated with this partition assignment (optional) */ - default void onAssignment(Assignment assignment, Optional generation) { + default void onAssignment(Assignment assignment, int generation) { onAssignment(assignment); } 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 3e67551eac1a1..a1fe0cd0e1cdb 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 @@ -652,7 +652,7 @@ public void testNoExceptionThrownWhenOnlySubscribedTopicDeleted() { Map> assignment = assignor.assign(partitionsPerTopic, subscriptions); subscriptions.put(consumer, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment( - new ConsumerUserData(assignment.get(consumer), 1)))); + new ConsumerUserData(assignment.get(consumer), Optional.of(1))))); assignment = assignor.assign(Collections.emptyMap(), subscriptions); assertEquals(assignment.size(), 1); diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index bf5027f57fd5a..3f3c1233da006 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -49,7 +49,6 @@ import java.util.LinkedList; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; @@ -794,11 +793,6 @@ List> interleaveTasksByGroupId(final Collection taskIds, fi return taskIdsForConsumerAssignment; } - @Override - public void onAssignment(final Assignment assignment, final Optional generation) { - onAssignment(assignment); - } - /** * @throws TaskAssignmentException if there is no task id for one of the partitions specified */ diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index 53f0fa0733482..284e396f978c6 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -53,7 +53,6 @@ import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; @@ -656,7 +655,7 @@ public void testOnAssignment() { EasyMock.expectLastCall(); EasyMock.replay(taskManager); - partitionAssignor.onAssignment(assignment, Optional.empty()); + partitionAssignor.onAssignment(assignment); EasyMock.verify(taskManager); @@ -977,7 +976,7 @@ public void shouldUpdateClusterMetadataAndHostInfoOnAssignment() { EasyMock.expectLastCall(); EasyMock.replay(taskManager); - partitionAssignor.onAssignment(createAssignment(hostState), Optional.empty()); + partitionAssignor.onAssignment(createAssignment(hostState)); EasyMock.verify(taskManager); } diff --git a/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java b/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java index 1d8e958be6a13..33e9b9771c245 100644 --- a/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java @@ -51,7 +51,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Properties; import java.util.Set; import java.util.UUID; @@ -142,9 +141,9 @@ public Subscription subscription(final Set topics) { } @Override - public void onAssignment(final PartitionAssignor.Assignment assignment, final Optional generation) { + public void onAssignment(final PartitionAssignor.Assignment assignment) { try { - super.onAssignment(assignment, generation); + super.onAssignment(assignment); return; } catch (final TaskAssignmentException cannotProcessFutureVersion) { // continue