You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by jg...@apache.org on 2021/10/19 23:16:46 UTC
[kafka] branch 2.7 updated: KAFKA-13081: detect and handle doubly
assigned partitions in cooperative assignor (#11068)
This is an automated email from the ASF dual-hosted git repository.
jgus pushed a commit to branch 2.7
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/2.7 by this push:
new 39af160 KAFKA-13081: detect and handle doubly assigned partitions in cooperative assignor (#11068)
39af160 is described below
commit 39af16054e8b26c9a6e0d53c3c149fec7abf1c43
Author: Luke Chen <sh...@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.7, 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 <ab...@apache.org>
---
checkstyle/suppressions.xml | 2 +-
.../consumer/CooperativeStickyAssignor.java | 44 ++++++++++-
.../consumer/internals/AbstractStickyAssignor.java | 65 ++++++++++++----
.../consumer/CooperativeStickyAssignorTest.java | 86 ++++++++++++++++++++--
.../kafka/clients/consumer/StickyAssignorTest.java | 48 ++++++++++--
.../internals/AbstractStickyAssignorTest.java | 65 ++++++++++++++++
6 files changed, 282 insertions(+), 28 deletions(-)
diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index 1665414..8b353b6 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -37,7 +37,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 aed8c09..9187014 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,24 @@
*/
package org.apache.kafka.clients.consumer;
-import static org.junit.Assert.assertTrue;
+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.Test;
+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 static java.util.Collections.emptyList;
+import static org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor.DEFAULT_GENERATION;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
public class CooperativeStickyAssignorTest extends AbstractStickyAssignorTest {
@@ -39,6 +47,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().intValue(), 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().intValue(), 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 2189700..71fb4f7 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.Assert.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 7ce765a..333a500 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
@@ -33,6 +33,7 @@ import org.apache.kafka.common.utils.Utils;
import org.junit.Before;
import org.junit.Test;
+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.Assert.assertEquals;
@@ -42,13 +43,21 @@ import static org.junit.Assert.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);
+
@Before
public void setUp() {
assignor = createAssignor();
@@ -649,6 +658,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);
}