Skip to content

KAFKA-12984: make AbstractStickyAssignor resilient to invalid input, utilize generation in cooperative, and fix assignment bug #10985

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
ableegoldman merged 9 commits into from
Jul 14, 2021
Merged

KAFKA-12984: make AbstractStickyAssignor resilient to invalid input, utilize generation in cooperative, and fix assignment bug #10985

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
24daac9
several fixes for better handling of invalid input and bug in assignm…
ableegoldman Jul 7, 2021
96c1378
clean up the fix, tests, checkstyle
ableegoldman Jul 7, 2021
e56913c
bump up timeout
ableegoldman Jul 7, 2021
2bcc470
encode generation with schema, add test
ableegoldman Jul 7, 2021
15acfa1
checkstyle
ableegoldman Jul 7, 2021
5f53140
review feedback
ableegoldman Jul 7, 2021
b32a644
log final assignment at INFO
ableegoldman Jul 8, 2021
3b66fa8
review feedback
ableegoldman Jul 13, 2021
e269e26
sort the unfilledMembersWithExactlyMinQuotaPartitions for determinist…
ableegoldman Jul 13, 2021
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
2 changes: 1 addition & 1 deletion checkstyle/suppressions.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"/>
Expand Down
30 changes: 29 additions & 1 deletion clients/src/main/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -43,6 +44,8 @@
*/
public class CooperativeStickyAssignor extends AbstractStickyAssignor {

private int generation = DEFAULT_GENERATION; // consumer group generation

@Override
public String name() {
return "cooperative-sticky";
Expand All @@ -53,9 +56,34 @@ public List<RebalanceProtocol> supportedProtocols() {
return Arrays.asList(RebalanceProtocol.COOPERATIVE, RebalanceProtocol.EAGER);
}

@Override
public void onAssignment(Assignment assignment, ConsumerGroupMetadata metadata) {
this.generation = metadata.generationId();
}

@Override
public ByteBuffer subscriptionUserData(Set<String> topics) {
if (generation == DEFAULT_GENERATION) {
Comment thread
ableegoldman marked this conversation as resolved.
Outdated
return ByteBuffer.allocate(0);
} else {
ByteBuffer buffer = ByteBuffer.allocate(4);
buffer.putInt(generation);
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 && buffer.rewind().hasRemaining()) {
encodedGeneration = Optional.of(buffer.getInt());
} else {
encodedGeneration = Optional.empty();
}
return new MemberData(subscription.ownedPartitions(), encodedGeneration);
}

@Override
Expand Down
116 changes: 86 additions & 30 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 @@ -28,6 +28,7 @@
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.Queue;
import java.util.Set;
import java.util.TreeSet;
import java.util.stream.Collectors;
Expand Down Expand Up @@ -71,32 +72,40 @@ public MemberData(List<TopicPartition> partitions, Optional<Integer> generation)
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 not all 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, consumerToOwnedPartitions);
}
}

/**
* 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) {
Map<String, List<TopicPartition>> consumerToOwnedPartitions,
Set<TopicPartition> partitionsWithMultiplePreviousOwners) {
Set<String> membersWithOldGeneration = new HashSet<>();
Set<String> membersOfCurrentHighestGeneration = new HashSet<>();
boolean isAllSubscriptionsEqual = true;

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();
Expand All @@ -122,6 +131,13 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
// 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();

@guozhangwang guozhangwang Jul 12, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

If the current member's generation is available but < maxGeneration, should we also clear it from the consumerToOwnedPartitions map? I think the passed in subscriptions is not sorted by the generations right?

@ableegoldman ableegoldman Jul 13, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

In that case, it's never added to consumerToOwnedPartitions in the first place. This map is not pre-filled, it gets populated inside this loop. So if its < maxGeneration, then we just insert an empty list into the map for that member's owned partitions

partitionsWithMultiplePreviousOwners.clear();
for (String droppedOutConsumer : membersWithOldGeneration) {
consumerToOwnedPartitions.get(droppedOutConsumer).clear();
}

@ableegoldman ableegoldman Jul 7, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

This part I just moved here to keep things up to date as we go, before we were clearing them after the loop

@showuon showuon Jul 7, 2021

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.

Because we cleared them earlier now, the membersWithOldGeneration is not necessary any more. We can just iterate membersOfCurrentHighestGeneration here.

@ableegoldman ableegoldman Jul 7, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Mm good catch


membersOfCurrentHighestGeneration.clear();
maxGeneration = memberData.generation.get();
}
Expand All @@ -130,19 +146,26 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
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.warn("Found multiple consumers {} and {} claiming the same TopicPartition {} in the "

@ableegoldman ableegoldman Jul 7, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

This is fix #2 -- if we somehow still get multiple consumers claiming a partition in the same generation, we have to consider both invalid and remove it from their ownedPartitions

@guozhangwang guozhangwang Jul 12, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

nit: do you think we should log at ERROR since this is not expected really? Right now we would sort of "hide" such bugs and still be able to proceed silently; I feel we should shouting out such scenarios a bit louder in logs.

@ableegoldman ableegoldman Jul 13, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Good point, yes I would absolutely want/hope a user would report this as a bug. Changed to ERROR

+ "same generation, this will be invalidated and removed from their previous assignment.",
consumer, otherConsumer, tp);
Comment thread
ableegoldman marked this conversation as resolved.
Outdated
consumerToOwnedPartitions.get(otherConsumer).remove(tp);
partitionsWithMultiplePreviousOwners.add(tp);
}
}
}
}
}

for (String consumer : membersWithOldGeneration) {
consumerToOwnedPartitions.get(consumer).clear();
}
return isAllSubscriptionsEqual;
}


/**
* This constrainedAssign optimizes the assignment algorithm when all consumers were subscribed to same set of topics.
* The method includes the following steps:
Expand All @@ -154,13 +177,15 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
* we're still under the number of expected max capacity members
* 2. Fill remaining members up to the expected numbers of maxQuota partitions, otherwise, to minQuota partitions
*
* @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.
*/
private Map<String, List<TopicPartition>> constrainedAssign(Map<String, Integer> partitionsPerTopic,
Map<String, List<TopicPartition>> consumerToOwnedPartitions) {
Map<String, List<TopicPartition>> consumerToOwnedPartitions,
Set<TopicPartition> partitionsWithMultiplePreviousOwners) {
if (log.isDebugEnabled()) {
log.debug("performing constrained assign. partitionsPerTopic: {}, consumerToOwnedPartitions: {}",
partitionsPerTopic, consumerToOwnedPartitions);
Expand All @@ -170,6 +195,7 @@ private Map<String, List<TopicPartition>> constrainedAssign(Map<String, Integer>

// the consumers not yet at expected capacity
List<String> unfilledMembers = new LinkedList<>();
Queue<String> potentiallyUnfilledMembersAtMinQuota = new LinkedList<>();

int numberOfConsumers = consumerToOwnedPartitions.size();
int totalPartitionsCount = partitionsPerTopic.values().stream().reduce(0, Integer::sum);
Expand All @@ -186,16 +212,25 @@ private Map<String, List<TopicPartition>> constrainedAssign(Map<String, Integer>
consumerToOwnedPartitions.keySet().stream().collect(Collectors.toMap(c -> c, c -> new ArrayList<>(maxQuota))));

List<TopicPartition> assignedPartitions = new ArrayList<>();
// Reassign previously owned partitions to the expected number
// Reassign previously owned partitions, up to the expected number of partitions per consumer
for (Map.Entry<String, List<TopicPartition>> consumerEntry : consumerToOwnedPartitions.entrySet()) {
String consumer = consumerEntry.getKey();
List<TopicPartition> ownedPartitions = consumerEntry.getValue();

List<TopicPartition> consumerAssignment = assignment.get(consumer);

for (TopicPartition doublyClaimedPartition : partitionsWithMultiplePreviousOwners) {
if (ownedPartitions.contains(doublyClaimedPartition)) {
log.warn("Found partition {} still claimed as owned by consumer {}, despite being claimed by multiple"

@ableegoldman ableegoldman Jul 7, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Strictly speaking this should never ever happen, even if we do get these "impossible" doubly-claimed partitions, we're also removing them from the ownedPartitions above (that's fix #2). But I put in a safeguard just in case, it shouldn't hurt (performance-wise we should generally not even enter this loop since partitionsWithMultiplePreviousOwners should almost always be empty

@showuon showuon Jul 7, 2021

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: add a space after "multiple". i.e. despite being claimed by multiple[ ]

+ "consumers already in the same generation. Removing it from the ownedPartitions",
doublyClaimedPartition, consumer);
ownedPartitions.remove(doublyClaimedPartition);
}
}

if (ownedPartitions.size() < minQuota) {
// the expected assignment size is more than consumer have now, so keep all the owned partitions
// and put this member into unfilled member list
// the expected assignment size is more than this consumer has now, so keep all the owned partitions
// and put this member into the unfilled member list
if (ownedPartitions.size() > 0) {
consumerAssignment.addAll(ownedPartitions);
assignedPartitions.addAll(ownedPartitions);
Expand All @@ -205,6 +240,9 @@ private Map<String, List<TopicPartition>> constrainedAssign(Map<String, Integer>
// consumer owned the "maxQuota" of partitions or more, and we're still under the number of expected members
// with more than the minQuota partitions, so keep "maxQuota" of the owned partitions, and revoke the rest of the partitions
numMembersAssignedOverMinQuota++;
if (numMembersAssignedOverMinQuota == expectedNumMembersAssignedOverMinQuota) {
potentiallyUnfilledMembersAtMinQuota.clear();

@guozhangwang guozhangwang Jul 12, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I'd suggest we remove this potentiallyUnfilledMembersAtMinQuota.clear(); logic and just keep two expected numbers:

  1. expectedNumMembersWithMaxQuota = totalPartitionsCount % numberOfConsumers;
  2. expectedNumMembersWithMinQuota = numberOfConsumers - expectedNumMembersWithMaxQuota;

And then we can also remove the check in line 309, and at the end when we exhausted unassignedPartitions, just check

numMembersWithMaxQuota == expectedNumMembersWithMaxQuota &&
numMembersWithMinQuota == expectedNumMembersWithMinQuota

WDYT?

@ableegoldman ableegoldman Jul 13, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

While I'm not really a fan of the potentiallyUnfilledMembersAtMinQuota logic (it's definitely awkward but I felt it was still the lesser evil in terms of complicating the code), I don't think we can get rid of it that easily. The problem is that when minQuota != maxQuota, and so far currentNumMembersWithOverMinQuotaPartitions < expectedNumMembersWithOverMinQuotaPartitions, then consumers that are filled up to exactly minQuota have to be considered potentially not yet at capacity since some will need one more partition, though not all. So this data structure is not just used to verify that everything is properly assigned after we've exhausted the unassignedPartitions, it's used to track which consumers can still receive another partition (ie, are "unfilled"). Does that make sense?

@guozhangwang guozhangwang Jul 13, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Yes that makes sense, still this logic

if (numMembersAssignedOverMinQuota == expectedNumMembersAssignedOverMinQuota) {
                    potentiallyUnfilledMembersAtMinQuota.clear();
}

Seems only needed because we have the check in 309 (?) Say if we do not check that, but instead just check the expected numbers of consumers with minQuota and maxQuota is satisfied, then do we still need this?

@showuon showuon Jul 13, 2021

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 logic Seems only needed because we have the check in 309 (?)

No, I don't think so. It should be for line 279:

// to handle the case that when there are still unassignedPartition left, but no more members to be assigned.
if (unfilledMembersWithUnderMinQuotaPartitions.isEmpty() && unfilledMembersWithExactlyMinQuotaPartitions.isEmpty()) {
   throw new IllegalStateException("No more unfilled consumers to be assigned.");

In line 309, it is just an early error detect and log for it. Not related to potentiallyUnfilledMembersAtMinQuota (or now unfilledMembersWithExactlyMinQuotaPartitions members)

@showuon showuon Jul 13, 2021

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 tried to summarize both methods:

What @guozhangwang 's meaning is, we can "lazily" detect the issue after assigning all unassignedPartitions. we don't need to clear the potentiallyUnfilledMembersAtMinQuota here, because as the "original" variable naming said: they are "potentially unfilled members", just keep them there. We "should not" assign partitions to them in this case because we've reached expectedNumMembersAssignedOverMinQuota.

But if somehow, after assigning unassignedPartitions to all unfilledMembers, there are still some unassignedPartitions left. We can just assign them to potentiallyUnfilledMembersAtMinQuota. And after running out the unassignedPartitions, we can check:

numMembersWithMaxQuota == expectedNumMembersWithMaxQuota &&
numMembersWithMinQuota == expectedNumMembersWithMinQuota

to do error handling.

VS.

In @ableegoldman 's version , we find issue immediately and handle it. we computed the potentiallyUnfilledMembersAtMinQuota correctly (that's why we need to clear it). So, if the issue happened:

if somehow, after assigning unassignedPartitions to all unfilledMembers, there are still some unassignedPartitions left

We can try to get member from potentiallyUnfilledMembersAtMinQuota and then assign unassignedPartition to the member. If we can't get member from it (i.e. potentiallyUnfilledMembersAtMinQuota is empty), we throw exception directly.

Both ways can find errors when happened. Personally, I like Sophie's version more since it's much clear.

@ableegoldman ableegoldman Jul 13, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Thanks @showuon and @guozhangwang , I think that all makes sense. One of my primary motivations was to keep all data structures at all times consistent with what they represent so they could always be relied upon to be used at any point. For that reason I also prefer to keep things as is, and clear the potentiallyUnfilledMembersAtMinQuota (now renamed to nfilledMembersWithExactlyMinQuotaPartitions) as soon as we have filled the last member who may be above minQuota, at which point all of the members at exactly minQuota are no longer considered "unfilled"

@guozhangwang guozhangwang Jul 13, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

SGTM.

}
List<TopicPartition> maxQuotaPartitions = ownedPartitions.subList(0, maxQuota);
consumerAssignment.addAll(maxQuotaPartitions);
assignedPartitions.addAll(maxQuotaPartitions);
Expand All @@ -218,15 +256,14 @@ private Map<String, List<TopicPartition>> constrainedAssign(Map<String, Integer>
allRevokedPartitions.addAll(ownedPartitions.subList(minQuota, ownedPartitions.size()));
// this consumer is potential maxQuota candidate since we're still under the number of expected members
// with more than the minQuota partitions. Note, if the number of expected members with more than
// the minQuota partitions is 0, it means minQuota == maxQuota, so they won't be put into unfilledMembers
// the minQuota partitions is 0, it means minQuota == maxQuota, and there are no potentially unfilled
if (numMembersAssignedOverMinQuota < expectedNumMembersAssignedOverMinQuota) {
unfilledMembers.add(consumer);
potentiallyUnfilledMembersAtMinQuota.add(consumer);

@ableegoldman ableegoldman Jul 7, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

This is part of fix #3 -- basically we have to handle these members separately, once they get up to minQuota they are only "potentially" unfilled. Once the last member allowed reaches maxQuota, all of these minQuota members are suddenly considered filled. cc @showuon

@guozhangwang guozhangwang Jul 12, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Honestly it took me quite a while to understand the fix :P After understanding that I think maybe it's better to rename these two collections more explicitly:

  1. unfilledMembers -> MembersWithLessThanMinQuotaPartitions.
  2. potentiallyUnfilledMembersAtMinQuota -> MembersWithExactMinQuotaPartitions.

And also (since the maxQuota is always either == minQuota or minQuota + 1):

  1. expectedNumMembersAssignedOverMinQuota -> expectedNumMembersWithMaxQuota
  2. numMembersAssignedOverMinQuota -> numMembersWithMaxQuota

@ableegoldman ableegoldman Jul 13, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Ack on the first two renamings, though I'd still want to prefix them with unfilled to emphasize that these structures only hold members that may potentially be assigned one or more partitions. ie, if minQuota == maxQuota, then potentiallyUnfilledMembersAtMinQuota should actually be empty, in which case MembersWithExactMinQuotaPartitions doesn't quite make sense. I'll clarify this in the comments as well.

@ableegoldman ableegoldman Jul 13, 2021 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

For the 3 & 4th suggested renamings, it's a bit subtle but this would actually be incorrect. In the case minQuota == maxQuota, the expectedNumMembersAssignedOverMinQuota variable will evaluate to 0, which would not make sense if it was called expectedNumMembersWithMaxQuota. Of course we could go through a tweak the logic for this case, but I'd prefer not to mix that into this PR. For now I'll just clarify in the comments for these variables.
(I did still rename them slightly to hopefully be more clear, and also in line with the new names of the other two data structures we renamed)

@guozhangwang guozhangwang Jul 13, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Ack.

}
}
}

List<TopicPartition> unassignedPartitions = getUnassignedPartitions(totalPartitionsCount, partitionsPerTopic, assignedPartitions);
assignedPartitions = null;

if (log.isDebugEnabled()) {
log.debug("After reassigning previously owned partitions, unfilled members: {}, unassigned partitions: {}, " +
Expand All @@ -238,32 +275,51 @@ private Map<String, List<TopicPartition>> constrainedAssign(Map<String, Integer>
Iterator<String> unfilledConsumerIter = unfilledMembers.iterator();
// Round-Robin filling remaining members up to the expected numbers of maxQuota, otherwise, to minQuota
for (TopicPartition unassignedPartition : unassignedPartitions) {
if (!unfilledConsumerIter.hasNext()) {
if (unfilledMembers.isEmpty()) {
// Should not enter here since we have calculated the exact number to assign to each consumer
// There might be issues in the assigning algorithm, or maybe assigning the same partition to two owners.
String consumer;
if (unfilledConsumerIter.hasNext()) {
consumer = unfilledConsumerIter.next();
} else {
if (unfilledMembers.isEmpty() && potentiallyUnfilledMembersAtMinQuota.isEmpty()) {
// Should not enter here since we have calculated the exact number to assign to each consumer.
// This indicates issues in the assignment algorithm
int currentPartitionIndex = unassignedPartitions.indexOf(unassignedPartition);
log.error("No more unfilled consumers to be assigned. The remaining unassigned partitions are: {}",
unassignedPartitions.subList(currentPartitionIndex, unassignedPartitions.size()));
unassignedPartitions.subList(currentPartitionIndex, unassignedPartitions.size()));
throw new IllegalStateException("No more unfilled consumers to be assigned.");
} else if (unfilledMembers.isEmpty()) {
consumer = potentiallyUnfilledMembersAtMinQuota.poll();
} else {
unfilledConsumerIter = unfilledMembers.iterator();
consumer = unfilledConsumerIter.next();
}
unfilledConsumerIter = unfilledMembers.iterator();
}
String consumer = unfilledConsumerIter.next();

List<TopicPartition> consumerAssignment = assignment.get(consumer);
consumerAssignment.add(unassignedPartition);

// We already assigned all possible ownedPartitions, so we know this must be newly assigned to this consumer
if (allRevokedPartitions.contains(unassignedPartition))
// or else the partition was actually claimed by multiple previous owners and had to be invalidated from all
// members claimed ownedPartitions
if (allRevokedPartitions.contains(unassignedPartition) || partitionsWithMultiplePreviousOwners.contains(unassignedPartition))
partitionsTransferringOwnership.put(unassignedPartition, consumer);
Comment thread
ableegoldman marked this conversation as resolved.

int currentAssignedCount = consumerAssignment.size();
int expectedAssignedCount = numMembersAssignedOverMinQuota < expectedNumMembersAssignedOverMinQuota ? maxQuota : minQuota;
Comment thread
ableegoldman marked this conversation as resolved.
Outdated
if (currentAssignedCount == expectedAssignedCount) {
if (currentAssignedCount == maxQuota) {
numMembersAssignedOverMinQuota++;
}
if (currentAssignedCount == minQuota) {
unfilledConsumerIter.remove();
potentiallyUnfilledMembersAtMinQuota.add(consumer);
} else if (currentAssignedCount == maxQuota) {
numMembersAssignedOverMinQuota++;
if (numMembersAssignedOverMinQuota == expectedNumMembersAssignedOverMinQuota) {
// We only start to iterate over the "potentially unfilled" members at minQuota after we've filled
// all members up to at least minQuota, so once the last minQuota member reaches maxQuota, we
// should be done. But in case of some algorithmic error, just log a warning and continue to
// assign any remaining partitions within the assignment constraints
if (unassignedPartitions.indexOf(unassignedPartition) != unassignedPartitions.size() - 1) {
log.warn("Filled the last member up to maxQuota but still had partitions remaining to assign, "

@guozhangwang guozhangwang Jul 12, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Related to the one above: maybe we just check that

numMembersWithMaxQuota == expectedNumMembersWithMaxQuota &&
numMembersWithMinQuota == expectedNumMembersWithMinQuota

And if not, log the full assignment as an ERROR?

@ableegoldman ableegoldman Jul 13, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I responded to the above comment as well, but specifically here I think that to just check on that condition requires us to make assumptions about the algorithm's correctness up to this point (and the correctness of its assumptions). But if those are all correct then we would never reach this to begin with, so it's better to directly look for any remaining unassignedPartitions -- it's a sanity check.

But ack on bumping to ERROR

+ "will continue but this indicates a bug in the assignment.");
}
}
}
}

Expand Down
Loading