kafka-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From ableegold...@apache.org
Subject [kafka] branch 2.8 updated: KAFKA-13081: detect and handle doubly assigned partitions in cooperative assignor (#11068)
Date Thu, 19 Aug 2021 00:04:44 GMT
This is an automated email from the ASF dual-hosted git repository.

ableegoldman pushed a commit to branch 2.8
in repository https://gitbox.apache.org/repos/asf/kafka.git


The following commit(s) were added to refs/heads/2.8 by this push:
     new 3d764ed  KAFKA-13081: detect and handle doubly assigned partitions in cooperative
assignor (#11068)
3d764ed is described below

commit 3d764ed38709aa0516e226afa8f4dac682bffb59
Author: Luke Chen <showuon@gmail.com>
AuthorDate: Thu Aug 19 08:03:22 2021 +0800

    KAFKA-13081: detect and handle doubly assigned partitions in cooperative assignor (#11068)
    
    This is the fix 1 and fix 2 in #10985 for v2.8, including the tests. Uses the generation
to invalidate previous assignments that claim partitions but no longer own them, and implements
an additional safety net to handle any case in which doubly-claimed partitions slip in to
the input anyway
    
    Reviewers: Anna Sophie Blee-Goldman <ableegoldman@apache.org>
---
 checkstyle/suppressions.xml                        |  2 +-
 .../consumer/CooperativeStickyAssignor.java        | 44 ++++++++++++-
 .../consumer/internals/AbstractStickyAssignor.java | 65 ++++++++++++++-----
 .../consumer/CooperativeStickyAssignorTest.java    | 75 ++++++++++++++++++++++
 .../kafka/clients/consumer/StickyAssignorTest.java | 48 ++++++++++++--
 .../internals/AbstractStickyAssignorTest.java      | 65 +++++++++++++++++++
 6 files changed, 276 insertions(+), 23 deletions(-)

diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index db760cb..8163f78 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -42,7 +42,7 @@
               files="AbstractResponse.java"/>
 
     <suppress checks="MethodLength"
-              files="KerberosLogin.java|RequestResponseTest.java|ConnectMetricsRegistry.java|KafkaConsumer.java"/>
+              files="(KerberosLogin|RequestResponseTest|ConnectMetricsRegistry|KafkaConsumer|AbstractStickyAssignor).java"/>
 
     <suppress checks="ParameterNumber"
               files="(NetworkClient|FieldSpec|KafkaRaftClient).java"/>
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignor.java
b/clients/src/main/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignor.java
index c7c0679..b2af7a0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignor.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -25,6 +26,10 @@ import java.util.Optional;
 import java.util.Set;
 import org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.protocol.types.Field;
+import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.protocol.types.Type;
 
 /**
  * A cooperative version of the {@link AbstractStickyAssignor AbstractStickyAssignor}. This
follows the same (sticky)
@@ -43,6 +48,15 @@ import org.apache.kafka.common.TopicPartition;
  */
 public class CooperativeStickyAssignor extends AbstractStickyAssignor {
 
+    // these schemas are used for preserving useful metadata for the assignment, such as
the last stable generation
+    private static final String GENERATION_KEY_NAME = "generation";
+    private static final Schema COOPERATIVE_STICKY_ASSIGNOR_USER_DATA_V0 = new Schema(
+        new Field(GENERATION_KEY_NAME, Type.INT32));
+
+    private int generation = DEFAULT_GENERATION; // consumer group generation
+
+
+
     @Override
     public String name() {
         return "cooperative-sticky";
@@ -54,8 +68,36 @@ public class CooperativeStickyAssignor extends AbstractStickyAssignor {
     }
 
     @Override
+    public void onAssignment(Assignment assignment, ConsumerGroupMetadata metadata) {
+        this.generation = metadata.generationId();
+    }
+
+    @Override
+    public ByteBuffer subscriptionUserData(Set<String> topics) {
+        Struct struct = new Struct(COOPERATIVE_STICKY_ASSIGNOR_USER_DATA_V0);
+
+        struct.set(GENERATION_KEY_NAME, generation);
+        ByteBuffer buffer = ByteBuffer.allocate(COOPERATIVE_STICKY_ASSIGNOR_USER_DATA_V0.sizeOf(struct));
+        COOPERATIVE_STICKY_ASSIGNOR_USER_DATA_V0.write(buffer, struct);
+        buffer.flip();
+        return buffer;
+    }
+
+    @Override
     protected MemberData memberData(Subscription subscription) {
-        return new MemberData(subscription.ownedPartitions(), Optional.empty());
+        ByteBuffer buffer = subscription.userData();
+        Optional<Integer> encodedGeneration;
+        if (buffer == null) {
+            encodedGeneration = Optional.empty();
+        } else {
+            try {
+                Struct struct = COOPERATIVE_STICKY_ASSIGNOR_USER_DATA_V0.read(buffer);
+                encodedGeneration = Optional.of(struct.getInt(GENERATION_KEY_NAME));
+            } catch (Exception e) {
+                encodedGeneration = Optional.of(DEFAULT_GENERATION);
+            }
+        }
+        return new MemberData(subscription.ownedPartitions(), encodedGeneration);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
index 7e42e44..3c5f609 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
@@ -73,14 +73,16 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor
{
     public Map<String, List<TopicPartition>> assign(Map<String, Integer>
partitionsPerTopic,
                                                     Map<String, Subscription> subscriptions)
{
         Map<String, List<TopicPartition>> consumerToOwnedPartitions = new HashMap<>();
-        if (allSubscriptionsEqual(partitionsPerTopic.keySet(), subscriptions, consumerToOwnedPartitions))
{
+        Set<TopicPartition> partitionsWithMultiplePreviousOwners = new HashSet<>();
+        if (allSubscriptionsEqual(partitionsPerTopic.keySet(), subscriptions, consumerToOwnedPartitions,
partitionsWithMultiplePreviousOwners)) {
             log.debug("Detected that all consumers were subscribed to same set of topics,
invoking the "
                           + "optimized assignment algorithm");
             partitionsTransferringOwnership = new HashMap<>();
-            return constrainedAssign(partitionsPerTopic, consumerToOwnedPartitions);
+            return constrainedAssign(partitionsPerTopic, consumerToOwnedPartitions, partitionsWithMultiplePreviousOwners);
         } else {
             log.debug("Detected that all not consumers were subscribed to same set of topics,
falling back to the "
                           + "general case assignment algorithm");
+            // we must set this to null for the general case so the cooperative assignor
knows to compute it from scratch
             partitionsTransferringOwnership = null;
             return generalAssign(partitionsPerTopic, subscriptions);
         }
@@ -88,17 +90,22 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor
{
 
     /**
      * Returns true iff all consumers have an identical subscription. Also fills out the
passed in
-     * {@code consumerToOwnedPartitions} with each consumer's previously owned and still-subscribed
partitions
+     * {@code consumerToOwnedPartitions} with each consumer's previously owned and still-subscribed
partitions,
+     * and the {@code partitionsWithMultiplePreviousOwners} with any partitions claimed by
multiple previous owners
      */
     private boolean allSubscriptionsEqual(Set<String> allTopics,
                                           Map<String, Subscription> subscriptions,
-                                          Map<String, List<TopicPartition>> consumerToOwnedPartitions)
{
-        Set<String> membersWithOldGeneration = new HashSet<>();
+                                          Map<String, List<TopicPartition>> consumerToOwnedPartitions,
+                                          Set<TopicPartition> partitionsWithMultiplePreviousOwners)
{
         Set<String> membersOfCurrentHighestGeneration = new HashSet<>();
         int maxGeneration = DEFAULT_GENERATION;
 
         Set<String> subscribedTopics = new HashSet<>();
 
+        // keep track of all previously owned partitions so we can invalidate them if invalid
input is
+        // detected, eg two consumers somehow claiming the same partition in the same/current
generation
+        Map<TopicPartition, String> allPreviousPartitionsToOwner = new HashMap<>();
+
         for (Map.Entry<String, Subscription> subscriptionEntry : subscriptions.entrySet())
{
             String consumer = subscriptionEntry.getKey();
             Subscription subscription = subscriptionEntry.getValue();
@@ -123,7 +130,12 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor
{
 
                 // If the current member's generation is higher, all the previously owned
partitions are invalid
                 if (memberData.generation.isPresent() && memberData.generation.get()
> maxGeneration) {
-                    membersWithOldGeneration.addAll(membersOfCurrentHighestGeneration);
+                    allPreviousPartitionsToOwner.clear();
+                    partitionsWithMultiplePreviousOwners.clear();
+                    for (String droppedOutConsumer : membersOfCurrentHighestGeneration) {
+                        consumerToOwnedPartitions.get(droppedOutConsumer).clear();
+                    }
+
                     membersOfCurrentHighestGeneration.clear();
                     maxGeneration = memberData.generation.get();
                 }
@@ -132,15 +144,22 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor
{
                 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())) {
-                        ownedPartitions.add(tp);
+                        if (!allPreviousPartitionsToOwner.containsKey(tp)) {
+                            allPreviousPartitionsToOwner.put(tp, consumer);
+                            ownedPartitions.add(tp);
+                        } else {
+                            String otherConsumer = allPreviousPartitionsToOwner.get(tp);
+                            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);
+                            partitionsWithMultiplePreviousOwners.add(tp);
+                        }
                     }
                 }
             }
         }
 
-        for (String consumer : membersWithOldGeneration) {
-            consumerToOwnedPartitions.get(consumer).clear();
-        }
         return true;
     }
 
@@ -156,13 +175,15 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor
{
      * 4. Otherwise we may have run out of unfilled consumers before assigning all partitions,
in which case we
      *    should just distribute one partition each to all consumers at min capacity
      *
-     * @param partitionsPerTopic          The number of partitions for each subscribed topic
-     * @param consumerToOwnedPartitions   Each consumer's previously owned and still-subscribed
partitions
+     * @param partitionsPerTopic                   The number of partitions for each subscribed
topic
+     * @param consumerToOwnedPartitions            Each consumer's previously owned and still-subscribed
partitions
+     * @param partitionsWithMultiplePreviousOwners The partitions being claimed in the previous
assignment of multiple consumers
      *
-     * @return Map from each member to the list of partitions assigned to them.
+     * @return                                     Map from each member to the list of partitions
assigned to them.
      */
     private Map<String, List<TopicPartition>> constrainedAssign(Map<String,
Integer> partitionsPerTopic,
-                                                                Map<String, List<TopicPartition>>
consumerToOwnedPartitions) {
+                                                                Map<String, List<TopicPartition>>
consumerToOwnedPartitions,
+                                                                Set<TopicPartition>
partitionsWithMultiplePreviousOwners) {
         SortedSet<TopicPartition> unassignedPartitions = getTopicPartitions(partitionsPerTopic);
 
         Set<TopicPartition> allRevokedPartitions = new HashSet<>();
@@ -189,6 +210,16 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor
{
             List<TopicPartition> ownedPartitions = consumerEntry.getValue();
 
             List<TopicPartition> consumerAssignment = assignment.get(consumer);
+
+            for (TopicPartition doublyClaimedPartition : partitionsWithMultiplePreviousOwners)
{
+                if (ownedPartitions.contains(doublyClaimedPartition)) {
+                    log.error("Found partition {} still claimed as owned by consumer {},
despite being claimed by multiple "
+                            + "consumers already in the same generation. Removing it from
the ownedPartitions",
+                        doublyClaimedPartition, consumer);
+                    ownedPartitions.remove(doublyClaimedPartition);
+                }
+            }
+
             int i = 0;
             // assign the first N partitions up to the max quota, and mark the remaining
as being revoked
             for (TopicPartition tp : ownedPartitions) {
@@ -228,7 +259,7 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor
{
                     consumerAssignment.add(tp);
                     unassignedPartitionsIter.remove();
                     // We already assigned all possible ownedPartitions, so we know this
must be newly to this consumer
-                    if (allRevokedPartitions.contains(tp))
+                    if (allRevokedPartitions.contains(tp) || partitionsWithMultiplePreviousOwners.contains(tp))
                         partitionsTransferringOwnership.put(tp, consumer);
                 } else {
                     break;
@@ -271,10 +302,12 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor
{
             // We can skip the bookkeeping of unassignedPartitions and maxCapacityMembers
here since we are at the end
             assignment.get(underCapacityConsumer).add(unassignedPartition);
 
-            if (allRevokedPartitions.contains(unassignedPartition))
+            if (allRevokedPartitions.contains(unassignedPartition) || partitionsWithMultiplePreviousOwners.contains(unassignedPartition))
                 partitionsTransferringOwnership.put(unassignedPartition, underCapacityConsumer);
         }
 
+        log.info("Final assignment of partitions to consumers: \n{}", assignment);
+
         return assignment;
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignorTest.java
b/clients/src/test/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignorTest.java
index d7d671e..29ee27a 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignorTest.java
@@ -16,16 +16,23 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import static java.util.Collections.emptyList;
+import static org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor.DEFAULT_GENERATION;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
 import org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor;
 import org.apache.kafka.clients.consumer.internals.AbstractStickyAssignorTest;
 import org.apache.kafka.common.TopicPartition;
+import org.junit.jupiter.api.Test;
 
 public class CooperativeStickyAssignorTest extends AbstractStickyAssignorTest {
 
@@ -39,6 +46,74 @@ public class CooperativeStickyAssignorTest extends AbstractStickyAssignorTest
{
         return new Subscription(topics, assignor.subscriptionUserData(new HashSet<>(topics)),
partitions);
     }
 
+    @Override
+    public Subscription buildSubscriptionWithGeneration(List<String> topics, List<TopicPartition>
partitions, int generation) {
+        assignor.onAssignment(null, new ConsumerGroupMetadata("dummy-group-id", generation,
"dummy-member-id", Optional.empty()));
+        return new Subscription(topics, assignor.subscriptionUserData(new HashSet<>(topics)),
partitions);
+    }
+
+    @Test
+    public void testEncodeAndDecodeGeneration() {
+        Subscription subscription = new Subscription(topics(topic), assignor.subscriptionUserData(new
HashSet<>(topics(topic))));
+
+        Optional<Integer> encodedGeneration = ((CooperativeStickyAssignor) assignor).memberData(subscription).generation;
+        assertTrue(encodedGeneration.isPresent());
+        assertEquals(encodedGeneration.get(), DEFAULT_GENERATION);
+
+        int generation = 10;
+        assignor.onAssignment(null, new ConsumerGroupMetadata("dummy-group-id", generation,
"dummy-member-id", Optional.empty()));
+
+        subscription = new Subscription(topics(topic), assignor.subscriptionUserData(new
HashSet<>(topics(topic))));
+        encodedGeneration = ((CooperativeStickyAssignor) assignor).memberData(subscription).generation;
+
+        assertTrue(encodedGeneration.isPresent());
+        assertEquals(encodedGeneration.get(), generation);
+    }
+
+    @Test
+    public void testDecodeGeneration() {
+        Subscription subscription = new Subscription(topics(topic));
+        assertFalse(((CooperativeStickyAssignor) assignor).memberData(subscription).generation.isPresent());
+    }
+
+    @Test
+    public void testAllConsumersHaveOwnedPartitionInvalidatedWhenClaimedByMultipleConsumersInSameGenerationWithEqualPartitionsPerConsumer()
{
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+
+        subscriptions.put(consumer1, buildSubscription(topics(topic), partitions(tp(topic,
0), tp(topic, 1))));
+        subscriptions.put(consumer2, buildSubscription(topics(topic), partitions(tp(topic,
0), tp(topic, 2))));
+        subscriptions.put(consumer3, buildSubscription(topics(topic), emptyList()));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
subscriptions);
+        assertEquals(partitions(tp(topic, 1)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic, 2)), assignment.get(consumer2));
+        // In the cooperative assignor, topic-0 has to be considered "owned" and so it cant
be assigned until both have "revoked" it
+        assertTrue(assignment.get(consumer3).isEmpty());
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testAllConsumersHaveOwnedPartitionInvalidatedWhenClaimedByMultipleConsumersInSameGenerationWithUnequalPartitionsPerConsumer()
{
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 4);
+
+        subscriptions.put(consumer1, buildSubscription(topics(topic), partitions(tp(topic,
0), tp(topic, 1))));
+        subscriptions.put(consumer2, buildSubscription(topics(topic), partitions(tp(topic,
0), tp(topic, 2))));
+        subscriptions.put(consumer3, buildSubscription(topics(topic), emptyList()));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
subscriptions);
+        assertEquals(partitions(tp(topic, 1)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic, 2), tp(topic, 3)), assignment.get(consumer2));
+        // In the cooperative assignor, topic-0 has to be considered "owned" and so it cant
be assigned until both have "revoked" it
+        assertTrue(assignment.get(consumer3).isEmpty());
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
     /**
      * The cooperative assignor must do some additional work and verification of some assignments
relative to the eager
      * assignor, since it may or may not need to trigger a second follow-up rebalance.
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 edc522b..678f62d 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
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import static java.util.Collections.emptyList;
 import static org.apache.kafka.clients.consumer.StickyAssignor.serializeTopicPartitionAssignment;
 import static org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor.DEFAULT_GENERATION;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -51,6 +52,48 @@ public class StickyAssignorTest extends AbstractStickyAssignorTest {
             serializeTopicPartitionAssignment(new MemberData(partitions, Optional.of(DEFAULT_GENERATION))));
     }
 
+    @Override
+    public Subscription buildSubscriptionWithGeneration(List<String> topics, List<TopicPartition>
partitions, int generation) {
+        return new Subscription(topics,
+            serializeTopicPartitionAssignment(new MemberData(partitions, Optional.of(generation))));
+    }
+
+    @Test
+    public void testAllConsumersHaveOwnedPartitionInvalidatedWhenClaimedByMultipleConsumersInSameGenerationWithEqualPartitionsPerConsumer()
{
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+
+        subscriptions.put(consumer1, buildSubscription(topics(topic), partitions(tp(topic,
0), tp(topic, 1))));
+        subscriptions.put(consumer2, buildSubscription(topics(topic), partitions(tp(topic,
0), tp(topic, 2))));
+        subscriptions.put(consumer3, buildSubscription(topics(topic), emptyList()));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
subscriptions);
+        assertEquals(partitions(tp(topic, 1)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic, 2)), assignment.get(consumer2));
+        assertEquals(partitions(tp(topic, 0)), assignment.get(consumer3));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testAllConsumersHaveOwnedPartitionInvalidatedWhenClaimedByMultipleConsumersInSameGenerationWithUnequalPartitionsPerConsumer()
{
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 4);
+
+        subscriptions.put(consumer1, buildSubscription(topics(topic), partitions(tp(topic,
0), tp(topic, 1))));
+        subscriptions.put(consumer2, buildSubscription(topics(topic), partitions(tp(topic,
0), tp(topic, 2))));
+        subscriptions.put(consumer3, buildSubscription(topics(topic), emptyList()));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
subscriptions);
+        assertEquals(partitions(tp(topic, 1)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic, 2), tp(topic, 3)), assignment.get(consumer2));
+        assertEquals(partitions(tp(topic, 0)), assignment.get(consumer3));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
     @Test
     public void testAssignmentWithMultipleGenerations1() {
         String consumer1 = "consumer1";
@@ -215,11 +258,6 @@ public class StickyAssignorTest extends AbstractStickyAssignorTest {
         assertTrue(isFullyBalanced(assignment));
     }
 
-    private Subscription buildSubscriptionWithGeneration(List<String> topics, List<TopicPartition>
partitions, int generation) {
-        return new Subscription(topics,
-            serializeTopicPartitionAssignment(new MemberData(partitions, Optional.of(generation))));
-    }
-
     private static Subscription buildSubscriptionWithOldSchema(List<String> topics,
List<TopicPartition> partitions) {
         Struct struct = new Struct(StickyAssignor.STICKY_ASSIGNOR_USER_DATA_V0);
         List<Struct> topicAssignments = new ArrayList<>();
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
index 3578540..c7b359d 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
@@ -34,6 +34,7 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
+import static org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor.DEFAULT_GENERATION;
 import static org.apache.kafka.common.utils.Utils.mkEntry;
 import static org.apache.kafka.common.utils.Utils.mkMap;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -43,13 +44,21 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
 public abstract class AbstractStickyAssignorTest {
     protected AbstractStickyAssignor assignor;
     protected String consumerId = "consumer";
+    protected String consumer1 = "consumer1";
+    protected String consumer2 = "consumer2";
+    protected String consumer3 = "consumer3";
     protected Map<String, Subscription> subscriptions;
     protected String topic = "topic";
+    protected String topic1 = "topic1";
+    protected String topic2 = "topic2";
+    protected String topic3 = "topic3";
 
     protected abstract AbstractStickyAssignor createAssignor();
 
     protected abstract Subscription buildSubscription(List<String> topics, List<TopicPartition>
partitions);
 
+    protected abstract Subscription buildSubscriptionWithGeneration(List<String> topics,
List<TopicPartition> partitions, int generation);
+
     @BeforeEach
     public void setUp() {
         assignor = createAssignor();
@@ -651,6 +660,62 @@ public abstract class AbstractStickyAssignorTest {
         }
     }
 
+    @Test
+    public void testAllConsumersReachExpectedQuotaAndAreConsideredFilled() {
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 4);
+
+        subscriptions.put(consumer1, buildSubscription(topics(topic), partitions(tp(topic,
0), tp(topic, 1))));
+        subscriptions.put(consumer2, buildSubscription(topics(topic), partitions(tp(topic,
2))));
+        subscriptions.put(consumer3, buildSubscription(topics(topic), Collections.emptyList()));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
subscriptions);
+        assertEquals(partitions(tp(topic, 0), tp(topic, 1)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic, 2)), assignment.get(consumer2));
+        assertEquals(partitions(tp(topic, 3)), assignment.get(consumer3));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testOwnedPartitionsAreInvalidatedForConsumerWithStaleGeneration() {
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+        partitionsPerTopic.put(topic2, 3);
+
+        int currentGeneration = 10;
+
+        subscriptions.put(consumer1, buildSubscriptionWithGeneration(topics(topic, topic2),
partitions(tp(topic, 0), tp(topic, 2), tp(topic2, 1)), currentGeneration));
+        subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic, topic2),
partitions(tp(topic, 0), tp(topic, 2), tp(topic2, 1)), currentGeneration - 1));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
subscriptions);
+        assertEquals(new HashSet<>(partitions(tp(topic, 0), tp(topic, 2), tp(topic2,
1))), new HashSet<>(assignment.get(consumer1)));
+        assertEquals(new HashSet<>(partitions(tp(topic, 1), tp(topic2, 0), tp(topic2,
2))), new HashSet<>(assignment.get(consumer2)));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testOwnedPartitionsAreInvalidatedForConsumerWithNoGeneration() {
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+        partitionsPerTopic.put(topic2, 3);
+
+        int currentGeneration = 10;
+
+        subscriptions.put(consumer1, buildSubscriptionWithGeneration(topics(topic, topic2),
partitions(tp(topic, 0), tp(topic, 2), tp(topic2, 1)), currentGeneration));
+        subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic, topic2),
partitions(tp(topic, 0), tp(topic, 2), tp(topic2, 1)), DEFAULT_GENERATION));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
subscriptions);
+        assertEquals(new HashSet<>(partitions(tp(topic, 0), tp(topic, 2), tp(topic2,
1))), new HashSet<>(assignment.get(consumer1)));
+        assertEquals(new HashSet<>(partitions(tp(topic, 1), tp(topic2, 0), tp(topic2,
2))), new HashSet<>(assignment.get(consumer2)));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
     private String getTopicName(int i, int maxNum) {
         return getCanonicalName("t", i, maxNum);
     }

Mime
View raw message