dajac commented on code in PR #13550:
URL: https://github.com/apache/kafka/pull/13550#discussion_r1178764451


##########
clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java:
##########
@@ -149,47 +148,57 @@ private boolean allSubscriptionsEqual(Set<String> 
allTopics,
             }
 
             MemberData memberData = memberData(subscription);
+            final int memberGeneration = 
memberData.generation.orElse(DEFAULT_GENERATION);
+            maxGeneration = Math.max(maxGeneration, memberGeneration);
 
             List<TopicPartition> ownedPartitions = new ArrayList<>();
             consumerToOwnedPartitions.put(consumer, ownedPartitions);
 
-            // Only consider this consumer's owned partitions as valid if it 
is a member of the current highest
-            // generation, or it's generation is not present but we have not 
seen any known generation so far
-            if (memberData.generation.isPresent() && 
memberData.generation.get() >= maxGeneration
-                || !memberData.generation.isPresent() && maxGeneration == 
DEFAULT_GENERATION) {
-
-                // If the current member's generation is higher, all the 
previously owned partitions are invalid
-                if (memberData.generation.isPresent() && 
memberData.generation.get() > maxGeneration) {
-                    allPreviousPartitionsToOwner.clear();
-                    partitionsWithMultiplePreviousOwners.clear();
-                    for (String droppedOutConsumer : 
membersOfCurrentHighestGeneration) {
-                        
consumerToOwnedPartitions.get(droppedOutConsumer).clear();
-                    }
-
-                    membersOfCurrentHighestGeneration.clear();
-                    maxGeneration = memberData.generation.get();
-                }
+            // the member has a valid generation, so we can consider its owned 
partitions if it has the highest
+            // generation amongst
+            for (final TopicPartition tp : memberData.partitions) {
+                if (allTopics.contains(tp.topic())) {
+                    String otherConsumer = 
allPreviousPartitionsToOwner.put(tp, consumer);
+                    if (otherConsumer == null) {
+                        // this partition is not owned by other consumer in 
the same generation
+                        ownedPartitions.add(tp);
+                    } else {
+                        final int otherMemberGeneration = 
subscriptions.get(otherConsumer).generationId().orElse(DEFAULT_GENERATION);
 
-                membersOfCurrentHighestGeneration.add(consumer);
-                for (final TopicPartition tp : memberData.partitions) {
-                    // filter out any topics that no longer exist or aren't 
part of the current subscription
-                    if (allTopics.contains(tp.topic())) {
-                        String otherConsumer = 
allPreviousPartitionsToOwner.put(tp, consumer);
-                        if (otherConsumer == null) {
-                            // this partition is not owned by other consumer 
in the same generation
-                            ownedPartitions.add(tp);
-                        } else {
+                        if (memberGeneration == otherMemberGeneration) {

Review Comment:
   nit: Could we put a comment in this branch like we did for the others?



##########
clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java:
##########
@@ -149,47 +148,57 @@ private boolean allSubscriptionsEqual(Set<String> 
allTopics,
             }
 
             MemberData memberData = memberData(subscription);
+            final int memberGeneration = 
memberData.generation.orElse(DEFAULT_GENERATION);
+            maxGeneration = Math.max(maxGeneration, memberGeneration);
 
             List<TopicPartition> ownedPartitions = new ArrayList<>();
             consumerToOwnedPartitions.put(consumer, ownedPartitions);
 
-            // Only consider this consumer's owned partitions as valid if it 
is a member of the current highest
-            // generation, or it's generation is not present but we have not 
seen any known generation so far
-            if (memberData.generation.isPresent() && 
memberData.generation.get() >= maxGeneration
-                || !memberData.generation.isPresent() && maxGeneration == 
DEFAULT_GENERATION) {
-
-                // If the current member's generation is higher, all the 
previously owned partitions are invalid
-                if (memberData.generation.isPresent() && 
memberData.generation.get() > maxGeneration) {
-                    allPreviousPartitionsToOwner.clear();
-                    partitionsWithMultiplePreviousOwners.clear();
-                    for (String droppedOutConsumer : 
membersOfCurrentHighestGeneration) {
-                        
consumerToOwnedPartitions.get(droppedOutConsumer).clear();
-                    }
-
-                    membersOfCurrentHighestGeneration.clear();
-                    maxGeneration = memberData.generation.get();
-                }
+            // the member has a valid generation, so we can consider its owned 
partitions if it has the highest
+            // generation amongst
+            for (final TopicPartition tp : memberData.partitions) {
+                if (allTopics.contains(tp.topic())) {
+                    String otherConsumer = 
allPreviousPartitionsToOwner.put(tp, consumer);
+                    if (otherConsumer == null) {
+                        // this partition is not owned by other consumer in 
the same generation
+                        ownedPartitions.add(tp);
+                    } else {
+                        final int otherMemberGeneration = 
subscriptions.get(otherConsumer).generationId().orElse(DEFAULT_GENERATION);
 
-                membersOfCurrentHighestGeneration.add(consumer);
-                for (final TopicPartition tp : memberData.partitions) {
-                    // filter out any topics that no longer exist or aren't 
part of the current subscription
-                    if (allTopics.contains(tp.topic())) {
-                        String otherConsumer = 
allPreviousPartitionsToOwner.put(tp, consumer);
-                        if (otherConsumer == null) {
-                            // this partition is not owned by other consumer 
in the same generation
-                            ownedPartitions.add(tp);
-                        } else {
+                        if (memberGeneration == otherMemberGeneration) {
                             log.error("Found multiple consumers {} and {} 
claiming the same TopicPartition {} in the "
-                                + "same generation {}, this will be 
invalidated and removed from their previous assignment.",
-                                     consumer, otherConsumer, tp, 
maxGeneration);
-                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                                            + "same generation {}, this will 
be invalidated and removed from their previous assignment.",
+                                    consumer, otherConsumer, tp, 
memberGeneration);
                             partitionsWithMultiplePreviousOwners.add(tp);
+                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                            allPreviousPartitionsToOwner.put(tp, consumer);
+                        } else if (memberGeneration > otherMemberGeneration) {
+                            // move partition from the member with an older 
generation to the member with the newer generation
+                            consumerToOwnedPartitions.get(consumer).add(tp);
+                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                            allPreviousPartitionsToOwner.put(tp, consumer);
+                            // if memberGeneration > otherMemberGeneration, 
the other member continue owns the generation
+                            log.warn("{} in generation {} and {} in generation 
{} claiming the same TopicPartition {} in " +

Review Comment:
   nit: I would say `Consumer {} in generation and consumer {} in ...`. This 
applies to the other comment as well.



##########
clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java:
##########
@@ -1038,6 +1038,99 @@ public void 
testPartitionsTransferringOwnershipIncludeThePartitionClaimedByMulti
         assertTrue(isFullyBalanced(assignment));
     }
 
+    @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
+    @EnumSource(RackConfig.class)
+    public void testEnsurePartitionsAssignedToHighestGeneration(RackConfig 
rackConfig) {
+        initializeRacks(rackConfig);
+        Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, partitionInfos(topic, 3));
+        partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
+        partitionsPerTopic.put(topic3, partitionInfos(topic3, 3));
+
+        int currentGeneration = 10;
+
+        // ensure partitions are always assigned to the member with the 
highest generation
+        // topic, 1 -> [consumer2], consumer3
+        // topic2, 1 -> [consumer2], consumer3
+        // topic, 0 -> [consumer1], consumer3
+        // topic3, 2 -> consumer3
+        subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic, 0), tp(topic2, 0), tp(topic3, 0)), 
currentGeneration, 0));
+        subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic, 1), tp(topic2, 1), tp(topic3, 1)), 
currentGeneration - 1, 1));
+        subscriptions.put(consumer3, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic3, 0), tp(topic3, 2), tp(topic2, 1)), 
currentGeneration - 2, 1));

Review Comment:
   nit: Could we put `tp(topic2, 1)` before `topic3`?



##########
clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java:
##########
@@ -1038,6 +1038,99 @@ public void 
testPartitionsTransferringOwnershipIncludeThePartitionClaimedByMulti
         assertTrue(isFullyBalanced(assignment));
     }
 
+    @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
+    @EnumSource(RackConfig.class)
+    public void testEnsurePartitionsAssignedToHighestGeneration(RackConfig 
rackConfig) {
+        initializeRacks(rackConfig);
+        Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, partitionInfos(topic, 3));
+        partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
+        partitionsPerTopic.put(topic3, partitionInfos(topic3, 3));
+
+        int currentGeneration = 10;
+
+        // ensure partitions are always assigned to the member with the 
highest generation
+        // topic, 1 -> [consumer2], consumer3
+        // topic2, 1 -> [consumer2], consumer3
+        // topic, 0 -> [consumer1], consumer3
+        // topic3, 2 -> consumer3
+        subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic, 0), tp(topic2, 0), tp(topic3, 0)), 
currentGeneration, 0));
+        subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic, 1), tp(topic2, 1), tp(topic3, 1)), 
currentGeneration - 1, 1));
+        subscriptions.put(consumer3, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic3, 0), tp(topic3, 2), tp(topic2, 1)), 
currentGeneration - 2, 1));
+
+        Map<String, List<TopicPartition>> assignment = 
assignor.assignPartitions(partitionsPerTopic, subscriptions);
+        assertEquals(new HashSet<>(partitions(tp(topic, 0), tp(topic2, 0), 
tp(topic3, 0))),
+            new HashSet<>(assignment.get(consumer1)));
+        assertEquals(new HashSet<>(partitions(tp(topic, 1), tp(topic2, 1), 
tp(topic3, 1))),
+            new HashSet<>(assignment.get(consumer2)));
+        assertEquals(new HashSet<>(partitions(tp(topic, 2), tp(topic2, 2), 
tp(topic3, 2))),
+            new HashSet<>(assignment.get(consumer3)));
+        assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
+
+        verifyValidityAndBalance(subscriptions, assignment, 
partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
+    @EnumSource(RackConfig.class)
+    public void testNoReassignmentOnCurrentMembers(RackConfig rackConfig) {
+        initializeRacks(rackConfig);
+        Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, partitionInfos(topic, 3));
+        partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
+        partitionsPerTopic.put(topic3, partitionInfos(topic3, 3));
+        partitionsPerTopic.put(topic1, partitionInfos(topic1, 3));

Review Comment:
   nit: Could we put `topic1` after `topic`?



##########
clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java:
##########
@@ -149,47 +148,57 @@ private boolean allSubscriptionsEqual(Set<String> 
allTopics,
             }
 
             MemberData memberData = memberData(subscription);
+            final int memberGeneration = 
memberData.generation.orElse(DEFAULT_GENERATION);
+            maxGeneration = Math.max(maxGeneration, memberGeneration);
 
             List<TopicPartition> ownedPartitions = new ArrayList<>();
             consumerToOwnedPartitions.put(consumer, ownedPartitions);
 
-            // Only consider this consumer's owned partitions as valid if it 
is a member of the current highest
-            // generation, or it's generation is not present but we have not 
seen any known generation so far
-            if (memberData.generation.isPresent() && 
memberData.generation.get() >= maxGeneration
-                || !memberData.generation.isPresent() && maxGeneration == 
DEFAULT_GENERATION) {
-
-                // If the current member's generation is higher, all the 
previously owned partitions are invalid
-                if (memberData.generation.isPresent() && 
memberData.generation.get() > maxGeneration) {
-                    allPreviousPartitionsToOwner.clear();
-                    partitionsWithMultiplePreviousOwners.clear();
-                    for (String droppedOutConsumer : 
membersOfCurrentHighestGeneration) {
-                        
consumerToOwnedPartitions.get(droppedOutConsumer).clear();
-                    }
-
-                    membersOfCurrentHighestGeneration.clear();
-                    maxGeneration = memberData.generation.get();
-                }
+            // the member has a valid generation, so we can consider its owned 
partitions if it has the highest
+            // generation amongst
+            for (final TopicPartition tp : memberData.partitions) {
+                if (allTopics.contains(tp.topic())) {
+                    String otherConsumer = 
allPreviousPartitionsToOwner.put(tp, consumer);
+                    if (otherConsumer == null) {
+                        // this partition is not owned by other consumer in 
the same generation
+                        ownedPartitions.add(tp);
+                    } else {
+                        final int otherMemberGeneration = 
subscriptions.get(otherConsumer).generationId().orElse(DEFAULT_GENERATION);
 
-                membersOfCurrentHighestGeneration.add(consumer);
-                for (final TopicPartition tp : memberData.partitions) {
-                    // filter out any topics that no longer exist or aren't 
part of the current subscription
-                    if (allTopics.contains(tp.topic())) {
-                        String otherConsumer = 
allPreviousPartitionsToOwner.put(tp, consumer);
-                        if (otherConsumer == null) {
-                            // this partition is not owned by other consumer 
in the same generation
-                            ownedPartitions.add(tp);
-                        } else {
+                        if (memberGeneration == otherMemberGeneration) {
                             log.error("Found multiple consumers {} and {} 
claiming the same TopicPartition {} in the "
-                                + "same generation {}, this will be 
invalidated and removed from their previous assignment.",
-                                     consumer, otherConsumer, tp, 
maxGeneration);
-                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                                            + "same generation {}, this will 
be invalidated and removed from their previous assignment.",
+                                    consumer, otherConsumer, tp, 
memberGeneration);
                             partitionsWithMultiplePreviousOwners.add(tp);
+                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                            allPreviousPartitionsToOwner.put(tp, consumer);
+                        } else if (memberGeneration > otherMemberGeneration) {
+                            // move partition from the member with an older 
generation to the member with the newer generation
+                            consumerToOwnedPartitions.get(consumer).add(tp);
+                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                            allPreviousPartitionsToOwner.put(tp, consumer);
+                            // if memberGeneration > otherMemberGeneration, 
the other member continue owns the generation
+                            log.warn("{} in generation {} and {} in generation 
{} claiming the same TopicPartition {} in " +
+                                            "different generations. The topic 
partition wil be assigned to the member with " +
+                                            "the higher generation {}.",
+                                    consumer, memberGeneration,
+                                    otherConsumer, otherMemberGeneration,
+                                    tp,
+                                    memberGeneration);
+                        } else {
+                            // if memberGeneration < otherMemberGeneration, 
the other member continue owns the generation

Review Comment:
   nit: Should we remove `if memberGeneration < otherMemberGeneration` as it is 
implicit?



##########
clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java:
##########
@@ -1038,6 +1038,99 @@ public void 
testPartitionsTransferringOwnershipIncludeThePartitionClaimedByMulti
         assertTrue(isFullyBalanced(assignment));
     }
 
+    @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
+    @EnumSource(RackConfig.class)
+    public void testEnsurePartitionsAssignedToHighestGeneration(RackConfig 
rackConfig) {
+        initializeRacks(rackConfig);
+        Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, partitionInfos(topic, 3));
+        partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
+        partitionsPerTopic.put(topic3, partitionInfos(topic3, 3));
+
+        int currentGeneration = 10;
+
+        // ensure partitions are always assigned to the member with the 
highest generation
+        // topic, 1 -> [consumer2], consumer3
+        // topic2, 1 -> [consumer2], consumer3
+        // topic, 0 -> [consumer1], consumer3
+        // topic3, 2 -> consumer3

Review Comment:
   I don't understand what you are trying to express here. It also seems to me 
that the mapping is incorrect. Should we just remove this?



##########
clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java:
##########
@@ -149,47 +148,57 @@ private boolean allSubscriptionsEqual(Set<String> 
allTopics,
             }
 
             MemberData memberData = memberData(subscription);
+            final int memberGeneration = 
memberData.generation.orElse(DEFAULT_GENERATION);
+            maxGeneration = Math.max(maxGeneration, memberGeneration);
 
             List<TopicPartition> ownedPartitions = new ArrayList<>();
             consumerToOwnedPartitions.put(consumer, ownedPartitions);
 
-            // Only consider this consumer's owned partitions as valid if it 
is a member of the current highest
-            // generation, or it's generation is not present but we have not 
seen any known generation so far
-            if (memberData.generation.isPresent() && 
memberData.generation.get() >= maxGeneration
-                || !memberData.generation.isPresent() && maxGeneration == 
DEFAULT_GENERATION) {
-
-                // If the current member's generation is higher, all the 
previously owned partitions are invalid
-                if (memberData.generation.isPresent() && 
memberData.generation.get() > maxGeneration) {
-                    allPreviousPartitionsToOwner.clear();
-                    partitionsWithMultiplePreviousOwners.clear();
-                    for (String droppedOutConsumer : 
membersOfCurrentHighestGeneration) {
-                        
consumerToOwnedPartitions.get(droppedOutConsumer).clear();
-                    }
-
-                    membersOfCurrentHighestGeneration.clear();
-                    maxGeneration = memberData.generation.get();
-                }
+            // the member has a valid generation, so we can consider its owned 
partitions if it has the highest
+            // generation amongst
+            for (final TopicPartition tp : memberData.partitions) {
+                if (allTopics.contains(tp.topic())) {
+                    String otherConsumer = 
allPreviousPartitionsToOwner.put(tp, consumer);
+                    if (otherConsumer == null) {
+                        // this partition is not owned by other consumer in 
the same generation
+                        ownedPartitions.add(tp);
+                    } else {
+                        final int otherMemberGeneration = 
subscriptions.get(otherConsumer).generationId().orElse(DEFAULT_GENERATION);
 
-                membersOfCurrentHighestGeneration.add(consumer);
-                for (final TopicPartition tp : memberData.partitions) {
-                    // filter out any topics that no longer exist or aren't 
part of the current subscription
-                    if (allTopics.contains(tp.topic())) {
-                        String otherConsumer = 
allPreviousPartitionsToOwner.put(tp, consumer);
-                        if (otherConsumer == null) {
-                            // this partition is not owned by other consumer 
in the same generation
-                            ownedPartitions.add(tp);
-                        } else {
+                        if (memberGeneration == otherMemberGeneration) {
                             log.error("Found multiple consumers {} and {} 
claiming the same TopicPartition {} in the "
-                                + "same generation {}, this will be 
invalidated and removed from their previous assignment.",
-                                     consumer, otherConsumer, tp, 
maxGeneration);
-                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                                            + "same generation {}, this will 
be invalidated and removed from their previous assignment.",
+                                    consumer, otherConsumer, tp, 
memberGeneration);
                             partitionsWithMultiplePreviousOwners.add(tp);
+                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                            allPreviousPartitionsToOwner.put(tp, consumer);
+                        } else if (memberGeneration > otherMemberGeneration) {
+                            // move partition from the member with an older 
generation to the member with the newer generation
+                            consumerToOwnedPartitions.get(consumer).add(tp);
+                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                            allPreviousPartitionsToOwner.put(tp, consumer);
+                            // if memberGeneration > otherMemberGeneration, 
the other member continue owns the generation

Review Comment:
   nit: Should we just remove this one? It seems that the comment earlier 
explain it all.



##########
clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java:
##########
@@ -1038,6 +1038,99 @@ public void 
testPartitionsTransferringOwnershipIncludeThePartitionClaimedByMulti
         assertTrue(isFullyBalanced(assignment));
     }
 
+    @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
+    @EnumSource(RackConfig.class)
+    public void testEnsurePartitionsAssignedToHighestGeneration(RackConfig 
rackConfig) {
+        initializeRacks(rackConfig);
+        Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, partitionInfos(topic, 3));
+        partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
+        partitionsPerTopic.put(topic3, partitionInfos(topic3, 3));
+
+        int currentGeneration = 10;
+
+        // ensure partitions are always assigned to the member with the 
highest generation
+        // topic, 1 -> [consumer2], consumer3
+        // topic2, 1 -> [consumer2], consumer3
+        // topic, 0 -> [consumer1], consumer3
+        // topic3, 2 -> consumer3
+        subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic, 0), tp(topic2, 0), tp(topic3, 0)), 
currentGeneration, 0));
+        subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic, 1), tp(topic2, 1), tp(topic3, 1)), 
currentGeneration - 1, 1));
+        subscriptions.put(consumer3, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic3, 0), tp(topic3, 2), tp(topic2, 1)), 
currentGeneration - 2, 1));
+
+        Map<String, List<TopicPartition>> assignment = 
assignor.assignPartitions(partitionsPerTopic, subscriptions);
+        assertEquals(new HashSet<>(partitions(tp(topic, 0), tp(topic2, 0), 
tp(topic3, 0))),
+            new HashSet<>(assignment.get(consumer1)));
+        assertEquals(new HashSet<>(partitions(tp(topic, 1), tp(topic2, 1), 
tp(topic3, 1))),
+            new HashSet<>(assignment.get(consumer2)));
+        assertEquals(new HashSet<>(partitions(tp(topic, 2), tp(topic2, 2), 
tp(topic3, 2))),
+            new HashSet<>(assignment.get(consumer3)));
+        assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
+
+        verifyValidityAndBalance(subscriptions, assignment, 
partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
+    @EnumSource(RackConfig.class)
+    public void testNoReassignmentOnCurrentMembers(RackConfig rackConfig) {
+        initializeRacks(rackConfig);
+        Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, partitionInfos(topic, 3));
+        partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
+        partitionsPerTopic.put(topic3, partitionInfos(topic3, 3));
+        partitionsPerTopic.put(topic1, partitionInfos(topic1, 3));
+
+        int currentGeneration = 10;
+
+        subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, 
topic2, topic3, topic1), partitions(),

Review Comment:
   nit: Could we put `partitions(),` on the next line to follow the pattern 
used by the others? It is easier to read and compare them.



##########
clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java:
##########
@@ -149,47 +148,57 @@ private boolean allSubscriptionsEqual(Set<String> 
allTopics,
             }
 
             MemberData memberData = memberData(subscription);
+            final int memberGeneration = 
memberData.generation.orElse(DEFAULT_GENERATION);
+            maxGeneration = Math.max(maxGeneration, memberGeneration);
 
             List<TopicPartition> ownedPartitions = new ArrayList<>();
             consumerToOwnedPartitions.put(consumer, ownedPartitions);
 
-            // Only consider this consumer's owned partitions as valid if it 
is a member of the current highest
-            // generation, or it's generation is not present but we have not 
seen any known generation so far
-            if (memberData.generation.isPresent() && 
memberData.generation.get() >= maxGeneration
-                || !memberData.generation.isPresent() && maxGeneration == 
DEFAULT_GENERATION) {
-
-                // If the current member's generation is higher, all the 
previously owned partitions are invalid
-                if (memberData.generation.isPresent() && 
memberData.generation.get() > maxGeneration) {
-                    allPreviousPartitionsToOwner.clear();
-                    partitionsWithMultiplePreviousOwners.clear();
-                    for (String droppedOutConsumer : 
membersOfCurrentHighestGeneration) {
-                        
consumerToOwnedPartitions.get(droppedOutConsumer).clear();
-                    }
-
-                    membersOfCurrentHighestGeneration.clear();
-                    maxGeneration = memberData.generation.get();
-                }
+            // the member has a valid generation, so we can consider its owned 
partitions if it has the highest
+            // generation amongst
+            for (final TopicPartition tp : memberData.partitions) {
+                if (allTopics.contains(tp.topic())) {
+                    String otherConsumer = 
allPreviousPartitionsToOwner.put(tp, consumer);
+                    if (otherConsumer == null) {
+                        // this partition is not owned by other consumer in 
the same generation
+                        ownedPartitions.add(tp);
+                    } else {
+                        final int otherMemberGeneration = 
subscriptions.get(otherConsumer).generationId().orElse(DEFAULT_GENERATION);
 
-                membersOfCurrentHighestGeneration.add(consumer);
-                for (final TopicPartition tp : memberData.partitions) {
-                    // filter out any topics that no longer exist or aren't 
part of the current subscription
-                    if (allTopics.contains(tp.topic())) {
-                        String otherConsumer = 
allPreviousPartitionsToOwner.put(tp, consumer);
-                        if (otherConsumer == null) {
-                            // this partition is not owned by other consumer 
in the same generation
-                            ownedPartitions.add(tp);
-                        } else {
+                        if (memberGeneration == otherMemberGeneration) {
                             log.error("Found multiple consumers {} and {} 
claiming the same TopicPartition {} in the "
-                                + "same generation {}, this will be 
invalidated and removed from their previous assignment.",
-                                     consumer, otherConsumer, tp, 
maxGeneration);
-                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                                            + "same generation {}, this will 
be invalidated and removed from their previous assignment.",
+                                    consumer, otherConsumer, tp, 
memberGeneration);
                             partitionsWithMultiplePreviousOwners.add(tp);
+                            
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
+                            allPreviousPartitionsToOwner.put(tp, consumer);
+                        } else if (memberGeneration > otherMemberGeneration) {
+                            // move partition from the member with an older 
generation to the member with the newer generation
+                            consumerToOwnedPartitions.get(consumer).add(tp);

Review Comment:
   nit: We could use `ownedPartitions.add(tp)`.



##########
clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java:
##########
@@ -1038,6 +1038,99 @@ public void 
testPartitionsTransferringOwnershipIncludeThePartitionClaimedByMulti
         assertTrue(isFullyBalanced(assignment));
     }
 
+    @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
+    @EnumSource(RackConfig.class)
+    public void testEnsurePartitionsAssignedToHighestGeneration(RackConfig 
rackConfig) {
+        initializeRacks(rackConfig);
+        Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, partitionInfos(topic, 3));
+        partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
+        partitionsPerTopic.put(topic3, partitionInfos(topic3, 3));
+
+        int currentGeneration = 10;
+
+        // ensure partitions are always assigned to the member with the 
highest generation
+        // topic, 1 -> [consumer2], consumer3
+        // topic2, 1 -> [consumer2], consumer3
+        // topic, 0 -> [consumer1], consumer3
+        // topic3, 2 -> consumer3
+        subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic, 0), tp(topic2, 0), tp(topic3, 0)), 
currentGeneration, 0));
+        subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic, 1), tp(topic2, 1), tp(topic3, 1)), 
currentGeneration - 1, 1));
+        subscriptions.put(consumer3, buildSubscriptionV2Above(topics(topic, 
topic2, topic3),
+            partitions(tp(topic3, 0), tp(topic3, 2), tp(topic2, 1)), 
currentGeneration - 2, 1));
+
+        Map<String, List<TopicPartition>> assignment = 
assignor.assignPartitions(partitionsPerTopic, subscriptions);
+        assertEquals(new HashSet<>(partitions(tp(topic, 0), tp(topic2, 0), 
tp(topic3, 0))),
+            new HashSet<>(assignment.get(consumer1)));
+        assertEquals(new HashSet<>(partitions(tp(topic, 1), tp(topic2, 1), 
tp(topic3, 1))),
+            new HashSet<>(assignment.get(consumer2)));
+        assertEquals(new HashSet<>(partitions(tp(topic, 2), tp(topic2, 2), 
tp(topic3, 2))),
+            new HashSet<>(assignment.get(consumer3)));
+        assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
+
+        verifyValidityAndBalance(subscriptions, assignment, 
partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
+    @EnumSource(RackConfig.class)
+    public void testNoReassignmentOnCurrentMembers(RackConfig rackConfig) {
+        initializeRacks(rackConfig);
+        Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, partitionInfos(topic, 3));
+        partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
+        partitionsPerTopic.put(topic3, partitionInfos(topic3, 3));
+        partitionsPerTopic.put(topic1, partitionInfos(topic1, 3));
+
+        int currentGeneration = 10;
+
+        subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, 
topic2, topic3, topic1), partitions(),
+            DEFAULT_GENERATION, 0));
+        subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, 
topic2, topic3, topic1),
+            partitions(tp(topic, 0), tp(topic2, 0), tp(topic1, 0)), 
currentGeneration - 1, 1));
+        subscriptions.put(consumer3, buildSubscriptionV2Above(topics(topic, 
topic2, topic3, topic1),
+            partitions(tp(topic3, 2), tp(topic2, 2), tp(topic1, 1)), 
currentGeneration - 2, 2));
+        subscriptions.put(consumer4, buildSubscriptionV2Above(topics(topic, 
topic2, topic3, topic1),
+            partitions(tp(topic3, 1), tp(topic, 1), tp(topic, 2)), 
currentGeneration - 3, 3));
+
+        Map<String, List<TopicPartition>> assignment = 
assignor.assignPartitions(partitionsPerTopic, subscriptions);
+        // ensure assigned partitions don't get reassigned
+        assertEquals(new HashSet<>(assignment.get(consumer1)),
+                new HashSet<>(partitions(tp(topic2, 1), tp(topic3, 0), 
tp(topic1, 2))));

Review Comment:
   nit: We should put the expected one first.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscr...@kafka.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Reply via email to