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 6 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
103 changes: 63 additions & 40 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 @@ -119,15 +120,13 @@ public Map<String, List<TopicPartition>> assign(Map<String, Integer> partitionsP
}

/**
* 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,
* and the {@code partitionsWithMultiplePreviousOwners} with any partitions claimed by multiple previous owners
* Returns the mapping of consumer to its owned TopicPartition, and fill 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<TopicPartition> partitionsWithMultiplePreviousOwners) {
Set<String> membersOfCurrentHighestGeneration = new HashSet<>();
boolean isAllSubscriptionsEqual = true;

Set<String> subscribedTopics = new HashSet<>();
Expand All @@ -137,7 +136,18 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
Map<TopicPartition, String> allPreviousPartitionsToOwner = new HashMap<>();

for (Map.Entry<String, Subscription> subscriptionEntry : subscriptions.entrySet()) {
String consumer = subscriptionEntry.getKey();
Subscription subscription = subscriptionEntry.getValue();
maxGeneration = Math.max(maxGeneration,
subscription.generationId().orElse(DEFAULT_GENERATION));
}

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()) {
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,46 +163,59 @@ 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
|| !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();
}
// if the consumer is earlier than the maxGeneration - 1, its partitions are irrelevant
if (!hasValidGeneration(maxGeneration - 1, memberData.generation) &&
!hasNoKnownGeneration(maxGeneration, memberData.generation)) {
continue;
}

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

if (memberGeneration == otherMemberGeneration) {

@dajac dajac Apr 27, 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 put a comment in this branch like we did for the others?

if (subscriptions.get(otherConsumer).generationId().orElse(DEFAULT_GENERATION) == memberData.generation.orElse(DEFAULT_GENERATION)) {

@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.

Aren't those two if statements the same?

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);

@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.

So my understanding is that partitions put in partitionsWithMultiplePreviousOwners will be unassigned from all consumers claiming them. Is my understanding correct?

@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.

that seems like the case, reference to the snippet here:

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);
                    }

@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.

From KAFKA-12984:

...the assignor will now explicitly look out for partitions that are being claimed by multiple consumers ... we have to invalidate this partition from the ownedPartitions of both consumers, since we can't tell who, if anyone, has the valid claim to this partition.

}
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
allPreviousPartitionsToOwner.put(tp, consumer);
continue;

@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.

This does not look nice. Should we remove it and use else if?

@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.

It could be. I got into the habit of returning early, I thought it makes it easier to read.

}

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);
if (memberGeneration > otherMemberGeneration) {
log.warn("Found multiple consumers {} and {} claiming the same TopicPartition {} in " +
"different " +
"generations. The topic partition wil be assigned to the member with higher generation.",
consumer, otherConsumer, tp);
consumerToOwnedPartitions.get(consumer).add(tp);

@dajac dajac Apr 27, 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: We could use ownedPartitions.add(tp).

// this partition is owned by other consumer in the same generation
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
partitionsWithMultiplePreviousOwners.add(tp);
allPreviousPartitionsToOwner.put(tp, consumer);
}
}
}
}
}

return isAllSubscriptionsEqual;
}

private static boolean hasValidGeneration(final int maxGeneration, final Optional<Integer> memberGeneration) {
return memberGeneration.orElse(Integer.MAX_VALUE) >= maxGeneration;
}

private static boolean hasNoKnownGeneration(final int maxGeneration, final Optional<Integer> memberGeneration) {
return !memberGeneration.isPresent() && maxGeneration == DEFAULT_GENERATION;
}

public boolean isSticky() {
return partitionMovements.isSticky();
}
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