Skip to content

KAFKA-14639: A single partition may be revoked and assign during a single round of rebalance #13550

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dajac merged 22 commits into from
Apr 28, 2023
Merged

KAFKA-14639: A single partition may be revoked and assign during a single round of rebalance #13550

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
76f0a15
Abstract sticky assignor should consider max and max -1 generation
Apr 12, 2023
93ad6f4
reset generation
Apr 12, 2023
08c4e94
Update AbstractStickyAssignor.java
Apr 12, 2023
23ae4e8
update doc
Apr 12, 2023
9610543
Merge branch 'apache:trunk' into kafka-13891-stickier-assignor
Apr 13, 2023
8ec312d
clean up
Apr 13, 2023
dbd0359
restore back to the original look
Apr 13, 2023
29bf846
Update AbstractStickyAssignor.java
Apr 13, 2023
6248852
Update AbstractStickyAssignor.java
Apr 13, 2023
3ac56a9
Merge branch 'apache:trunk' into kafka-13891-stickier-assignor
Apr 17, 2023
2c468b1
resend owned partitions
Apr 18, 2023
ee10026
doesn't need to be abstract
Apr 18, 2023
152c102
refactor the tests
Apr 18, 2023
319886f
clean up
Apr 19, 2023
946a31e
handling eager case, and retain the last generation
Apr 19, 2023
1d3deb3
clean up
Apr 19, 2023
5ed02a7
un-optimize import and clean up
Apr 19, 2023
5c93207
assignor takes the highest gen partition owner
Apr 25, 2023
897321f
member with the highest generation should maintain its partitions
Apr 25, 2023
9cfefad
ensure no reassignment
Apr 25, 2023
767c04d
PR comments and test improvements
Apr 26, 2023
cf3d778
fix some nits
Apr 27, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -835,6 +835,7 @@ public void handle(SyncGroupResponse syncResponse,
} else if (error == Errors.REBALANCE_IN_PROGRESS) {
log.info("SyncGroup failed: The group began another rebalance. Need to re-join the group. " +
"Sent generation was {}", sentGeneration);
resetStateAndGeneration("member missed the rebalance", true);
Comment thread
philipnee marked this conversation as resolved.
Outdated
future.raise(error);
} else if (error == Errors.FENCED_INSTANCE_ID) {
// for sync-group request, even if the generation has changed we would not expect the instance id
Expand Down
69 changes: 38 additions & 31 deletions ...nts/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,13 @@
*/
package org.apache.kafka.clients.consumer.internals;

import org.apache.kafka.clients.consumer.internals.Utils.PartitionComparator;

@philipnee philipnee Apr 13, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is "optimized" by the IDE

import org.apache.kafka.common.Node;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
Expand All @@ -33,12 +40,6 @@
import java.util.Set;
import java.util.TreeSet;
import java.util.stream.Collectors;
import org.apache.kafka.clients.consumer.internals.Utils.PartitionComparator;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
* Sticky assignment implementation used by {@link org.apache.kafka.clients.consumer.StickyAssignor} and
Expand Down Expand Up @@ -127,7 +128,6 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
Map<String, Subscription> subscriptions,
Map<String, List<TopicPartition>> consumerToOwnedPartitions,
Set<TopicPartition> partitionsWithMultiplePreviousOwners) {
Set<String> membersOfCurrentHighestGeneration = new HashSet<>();
boolean isAllSubscriptionsEqual = true;

Set<String> subscribedTopics = new HashSet<>();
Expand All @@ -136,8 +136,13 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
// detected, eg two consumers somehow claiming the same partition in the same/current generation
Map<TopicPartition, String> allPreviousPartitionsToOwner = new HashMap<>();

maxGeneration = subscriptions.values().stream()
.map(v -> v.generationId().orElse(DEFAULT_GENERATION))
.max(Integer::compareTo)
.orElse(DEFAULT_GENERATION);

for (Map.Entry<String, Subscription> subscriptionEntry : subscriptions.entrySet()) {
String consumer = subscriptionEntry.getKey();
final String consumer = subscriptionEntry.getKey();
Subscription subscription = subscriptionEntry.getValue();

@dajac dajac Apr 26, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: While here, should this one be final as well?


// initialize the subscribed topics set if this is the first subscription
Expand All @@ -153,43 +158,45 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
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
if (memberData.generation.isPresent() && memberData.generation.get() >= maxGeneration - 1
|| !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) {
Comment thread
philipnee marked this conversation as resolved.
allPreviousPartitionsToOwner.clear();
partitionsWithMultiplePreviousOwners.clear();
for (String droppedOutConsumer : membersOfCurrentHighestGeneration) {
consumerToOwnedPartitions.get(droppedOutConsumer).clear();
}

membersOfCurrentHighestGeneration.clear();
maxGeneration = memberData.generation.get();
}

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 {
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);
int memberGeneration = memberData.generation.orElse(DEFAULT_GENERATION);
int otherMemberGeneration = subscriptions.get(otherConsumer).generationId().orElse(DEFAULT_GENERATION);

if (memberGeneration == otherMemberGeneration) {
if (subscriptions.get(otherConsumer).generationId().orElse(DEFAULT_GENERATION) == memberData.generation.orElse(DEFAULT_GENERATION)) {
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);
partitionsWithMultiplePreviousOwners.add(tp);
}
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
allPreviousPartitionsToOwner.put(tp, consumer);
continue;
}

if (memberGeneration > otherMemberGeneration) {
log.warn("Found multiple consumers {} and {} claiming the same TopicPartition {} in " +
Comment thread
philipnee marked this conversation as resolved.
Outdated
"different generations. The topic partition wil be assigned to the member with higher generation.",
consumer, otherConsumer, tp);
// 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);
}
}
}
}
}
}

return isAllSubscriptionsEqual;
}

Expand Down
40 changes: 40 additions & 0 deletions ...ts/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -549,6 +549,46 @@ public void testRejoinReason() {
assertEquals("", coordinator.rejoinReason());
}

@Test
public void testSyncGroupRebalanceInProgressResetGeneration() throws InterruptedException {
setupCoordinator();
joinGroup();

final AbstractCoordinator.Generation generation = coordinator.generation();

coordinator.setNewState(AbstractCoordinator.MemberState.PREPARING_REBALANCE);
RequestFuture<ByteBuffer> future = coordinator.sendJoinGroupRequest();

// successfully join group
TestUtils.waitForCondition(() -> {
consumerClient.poll(mockTime.timer(REQUEST_TIMEOUT_MS));
return !mockClient.requests().isEmpty();
}, 2000, "The join-group request was not sent");
mockClient.respond(joinGroupFollowerResponse(generation.generationId,
memberId,
JoinGroupRequest.UNKNOWN_MEMBER_ID,
Errors.NONE));
assertTrue(mockClient.requests().isEmpty());

TestUtils.waitForCondition(() -> {
consumerClient.poll(mockTime.timer(REQUEST_TIMEOUT_MS));
return !mockClient.requests().isEmpty();
}, 2000, "The sync-group request was not sent");

// sync group response with REBALANCE_IN_PROGRESS should reset the generation
final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
generation.generationId,
generation.memberId + "-new",
generation.protocolName);
coordinator.setNewGeneration(newGen);

mockClient.respond(syncGroupResponse(Errors.REBALANCE_IN_PROGRESS));
assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
assertTrue(future.exception().getClass().isInstance(Errors.REBALANCE_IN_PROGRESS.exception()));

assertEquals(AbstractCoordinator.Generation.NO_GENERATION, coordinator.generation());
}

private void ensureActiveGroup(
int generation,
String memberId
Expand Down
83 changes: 70 additions & 13 deletions ...src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,18 +16,6 @@
*/
package org.apache.kafka.clients.consumer.internals;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
import java.util.Set;

import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
import org.apache.kafka.clients.consumer.StickyAssignor;
import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.RackConfig;
Expand All @@ -42,10 +30,22 @@
import org.junit.jupiter.params.provider.EnumSource;
import org.junit.jupiter.params.provider.ValueSource;

import java.nio.ByteBuffer;

@dajac dajac Apr 26, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Could we revert this as it is not related to the fix?

@philipnee philipnee Apr 26, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh yes, will do it.... thanks for my IDE's import optimization.

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
import java.util.Set;

import static java.util.Arrays.asList;
import static java.util.Collections.emptyList;
import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.TEST_NAME_WITH_RACK_CONFIG;
import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.TEST_NAME_WITH_CONSUMER_RACK;
import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.TEST_NAME_WITH_RACK_CONFIG;
import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.nullRacks;
import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.racks;
import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.verifyRackAssignment;
Expand Down Expand Up @@ -1038,6 +1038,63 @@ public void testPartitionsTransferringOwnershipIncludeThePartitionClaimedByMulti
assertTrue(isFullyBalanced(assignment));
}

@ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
@EnumSource(RackConfig.class)
public void testOwnedPartitionsAreStableForConsumerWithMultipleGeneration(RackConfig rackConfig) {

@dajac dajac Apr 26, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: It seems that we also test the case where a partition is claimed by two consumers with different generation in this test. Should we try to reflect this in its name?

@dajac dajac Apr 26, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I actually wonder if this was expected or if it is just a typo in the test case because there is also testOwnedPartitionsAreInvalidatedForConsumerWithMultipleGeneration.

@philipnee philipnee Apr 26, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok I renamed the test to testEnsurePartitionsAssignedToHighestGeneration as the goal of this test is to make sure partitions are always assigned to the member with the highest generation.

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;

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(topic, 2), tp(topic2, 2), tp(topic3, 2)), 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 testOwnedPartitionsAreInvalidatedForConsumerWithMultipleGeneration(RackConfig rackConfig) {
initializeRacks(rackConfig);
Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
partitionsPerTopic.put(topic, partitionInfos(topic, 3));
partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));

int currentGeneration = 10;

subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, topic2),
partitions(tp(topic, 0), tp(topic2, 1), tp(topic, 1)), currentGeneration, 0));
subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, topic2),
partitions(tp(topic, 0), tp(topic2, 1), tp(topic2, 2)), currentGeneration - 2, 1));

Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
assertEquals(new HashSet<>(partitions(tp(topic, 0), tp(topic2, 1), tp(topic, 1))),
new HashSet<>(assignment.get(consumer1)));
assertEquals(new HashSet<>(partitions(tp(topic, 2), tp(topic2, 2), tp(topic2, 0))),
new HashSet<>(assignment.get(consumer2)));
assertTrue(assignor.partitionsTransferringOwnership.isEmpty());

verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
assertTrue(isFullyBalanced(assignment));
}

@dajac dajac Apr 26, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we also add a test for the case where a partition is claimed by two consumers with the same generation?

@philipnee philipnee Apr 26, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good idea, I didn't check because i thought there could be an existing test case covering it but let me check

@philipnee philipnee Apr 26, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok testPartitionsTransferringOwnershipIncludeThePartitionClaimedByMultipleConsumersInSameGeneration might cover it. But I'll double check if that's what we want.

@Test
public void testRackAwareAssignmentWithUniformSubscription() {
Map<String, Integer> topics = mkMap(mkEntry("t1", 6), mkEntry("t2", 7), mkEntry("t3", 2));
Expand Down