You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by ab...@apache.org on 2024/05/15 00:12:28 UTC
(kafka) 02/02: KAFKA-15170: Fix rack-aware assignment in AbstractStickyAssignor (#13965)
This is an automated email from the ASF dual-hosted git repository.
ableegoldman pushed a commit to branch 3.7
in repository https://gitbox.apache.org/repos/asf/kafka.git
commit 2eb6baac1d211de2cf734347441f6730e6b80a08
Author: flashmouse <li...@tuya.com>
AuthorDate: Wed May 15 08:06:05 2024 +0800
KAFKA-15170: Fix rack-aware assignment in AbstractStickyAssignor (#13965)
This PR fixes two kinds of bugs in the new(ish) rack-aware part of the sticky assignment algorithm:
First, when reassigning "owned partitions" to their previous owners, we now have to take the rack placement into account and might not immediately assign a previously-owned partition to its old consumer during this phase. There is a small chance this partition will be assigned to its previous owner during a later stage of the assignment, but if it's not then by definition it has been "revoked" and must be removed from the assignment during the adjustment phase of the CooperativeSticky [...]
Second, the sticky algorithm works in part by keeping track of how many consumers are still "unfilled" when they are at the "minQuota", meaning we may need to assign one more partition to get to the expected number of consumers at the "maxQuota". During the rack-aware round-robin assignment phase, we were not properly clearing the set of unfilled & minQuota consumers once we reached the expected number of "maxQuota" consumers (since by definition that means no more minQuota consumers [...]
Reviewers: Anna Sophie Blee-Goldman <ab...@apache.org>
---
.../consumer/internals/AbstractStickyAssignor.java | 43 ++++++++++++++++++----
.../consumer/CooperativeStickyAssignorTest.java | 17 +++++++++
.../internals/AbstractStickyAssignorTest.java | 27 ++++++++++----
3 files changed, 72 insertions(+), 15 deletions(-)
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 727dac231ac..587b05968b2 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
@@ -211,6 +211,10 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor {
return partitionMovements.isSticky();
}
+ public Map<TopicPartition, String> partitionsTransferringOwnership() {
+ return partitionsTransferringOwnership;
+ }
+
private static class TopicComparator implements Comparator<String>, Serializable {
private static final long serialVersionUID = 1L;
private final Map<String, List<String>> map;
@@ -579,7 +583,7 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor {
private class ConstrainedAssignmentBuilder extends AbstractAssignmentBuilder {
private final Set<TopicPartition> partitionsWithMultiplePreviousOwners;
- private final Set<TopicPartition> allRevokedPartitions;
+ private final Map<TopicPartition, String> maybeRevokedPartitions;
// the consumers which may still be assigned one or more partitions to reach expected capacity
private final List<String> unfilledMembersWithUnderMinQuotaPartitions;
@@ -610,7 +614,7 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor {
super(partitionsPerTopic, rackInfo, consumerToOwnedPartitions);
this.partitionsWithMultiplePreviousOwners = partitionsWithMultiplePreviousOwners;
- allRevokedPartitions = new HashSet<>();
+ maybeRevokedPartitions = new HashMap<>();
unfilledMembersWithUnderMinQuotaPartitions = new LinkedList<>();
unfilledMembersWithExactlyMinQuotaPartitions = new LinkedList<>();
@@ -662,7 +666,13 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor {
for (Map.Entry<String, List<TopicPartition>> consumerEntry : currentAssignment.entrySet()) {
String consumer = consumerEntry.getKey();
List<TopicPartition> ownedPartitions = consumerEntry.getValue().stream()
- .filter(tp -> !rackInfo.racksMismatch(consumer, tp))
+ .filter(tp -> {
+ boolean mismatch = rackInfo.racksMismatch(consumer, tp);
+ if (mismatch) {
+ maybeRevokedPartitions.put(tp, consumer);
+ }
+ return !mismatch;
+ })
.sorted(Comparator.comparing(TopicPartition::partition).thenComparing(TopicPartition::topic))
.collect(Collectors.toList());
@@ -695,14 +705,18 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor {
List<TopicPartition> maxQuotaPartitions = ownedPartitions.subList(0, maxQuota);
consumerAssignment.addAll(maxQuotaPartitions);
assignedPartitions.addAll(maxQuotaPartitions);
- allRevokedPartitions.addAll(ownedPartitions.subList(maxQuota, ownedPartitions.size()));
+ for (TopicPartition topicPartition : ownedPartitions.subList(maxQuota, ownedPartitions.size())) {
+ maybeRevokedPartitions.put(topicPartition, consumer);
+ }
} else {
// consumer owned at least "minQuota" of partitions
// so keep "minQuota" of the owned partitions, and revoke the rest of the partitions
List<TopicPartition> minQuotaPartitions = ownedPartitions.subList(0, minQuota);
consumerAssignment.addAll(minQuotaPartitions);
assignedPartitions.addAll(minQuotaPartitions);
- allRevokedPartitions.addAll(ownedPartitions.subList(minQuota, ownedPartitions.size()));
+ for (TopicPartition topicPartition : ownedPartitions.subList(minQuota, ownedPartitions.size())) {
+ maybeRevokedPartitions.put(topicPartition, consumer);
+ }
// 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, and there are no potentially unfilled
@@ -729,8 +743,12 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor {
int assignmentCount = assignment.get(consumer).size() + 1;
if (assignmentCount >= minQuota) {
unfilledMembersWithUnderMinQuotaPartitions.remove(consumer);
- if (assignmentCount < maxQuota)
+ // Only add this consumer if the current num members at maxQuota is less than the expected number
+ // since a consumer at minQuota can only be considered unfilled if it's possible to add another partition,
+ // which would bump it to maxQuota and exceed the expectedNumMembersWithOverMinQuotaPartitions
+ if (assignmentCount < maxQuota && (currentNumMembersWithOverMinQuotaPartitions < expectedNumMembersWithOverMinQuotaPartitions)) {
unfilledMembersWithExactlyMinQuotaPartitions.add(consumer);
+ }
} else {
nextIndex++;
}
@@ -739,8 +757,15 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor {
int firstIndex = rackInfo.nextRackConsumer(unassignedPartition, unfilledMembersWithExactlyMinQuotaPartitions, 0);
if (firstIndex >= 0) {
consumer = unfilledMembersWithExactlyMinQuotaPartitions.get(firstIndex);
- if (assignment.get(consumer).size() + 1 == maxQuota)
+ if (assignment.get(consumer).size() + 1 == maxQuota) {
unfilledMembersWithExactlyMinQuotaPartitions.remove(firstIndex);
+ currentNumMembersWithOverMinQuotaPartitions++;
+ // Clear this once the current num consumers over minQuota reaches the expected number since this
+ // means all consumers at minQuota are now considered filled
+ if (currentNumMembersWithOverMinQuotaPartitions == expectedNumMembersWithOverMinQuotaPartitions) {
+ unfilledMembersWithExactlyMinQuotaPartitions.clear();
+ }
+ }
}
}
@@ -803,8 +828,10 @@ public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor {
// We already assigned all possible ownedPartitions, so we know this must be newly assigned to this consumer
// 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))
+ if ((maybeRevokedPartitions.containsKey(unassignedPartition) && !maybeRevokedPartitions.get(unassignedPartition).equals(consumer))
+ || partitionsWithMultiplePreviousOwners.contains(unassignedPartition)) {
partitionsTransferringOwnership.put(unassignedPartition, consumer);
+ }
return consumerAssignment.size();
}
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 65f049922f5..23c8836395d 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
@@ -207,6 +207,23 @@ public class CooperativeStickyAssignorTest extends AbstractStickyAssignorTest {
assertTrue(isFullyBalanced(assignment));
}
+ @Test
+ public void testUniformSubscriptionTransferOwnershipListIsRight() {
+ this.replicationFactor = 1;
+ this.numBrokerRacks = 2;
+ this.hasConsumerRack = true;
+ Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
+ partitionsPerTopic.put(topic1, partitionInfos(topic1, 4));
+
+ subscriptions.put("c0", buildSubscriptionV2Above(topics(topic1), partitions(tp(topic1, 0), tp(topic1, 1)),
+ generationId, 0));
+ subscriptions.put("c1", buildSubscriptionV2Above(topics(topic1), partitions(tp(topic1, 2), tp(topic1, 3)),
+ generationId, 1));
+
+ assignor.assignPartitions(partitionsPerTopic, subscriptions);
+ assertEquals(2, assignor.partitionsTransferringOwnership().size());
+ }
+
/**
* 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/internals/AbstractStickyAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
index 1b771929c4d..7fe35977e45 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
@@ -30,6 +30,7 @@ import java.util.Set;
import java.util.stream.Collectors;
import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
+import org.apache.kafka.clients.consumer.CooperativeStickyAssignor;
import org.apache.kafka.clients.consumer.StickyAssignor;
import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.RackConfig;
import org.apache.kafka.common.PartitionInfo;
@@ -1297,14 +1298,26 @@ public abstract class AbstractStickyAssignorTest {
assignment = asList("t1-0, t1-3, t2-0, t2-3, t2-6", "t1-1, t1-4, t2-1, t2-4, t3-0", "t1-2, t1-5, t2-2, t2-5, t3-1");
verifyUniformSubscription(assignor, topics, 3, racks(2), racks(3), consumerTopics, assignment, 5);
- // Verify that rack-awareness is improved if already owned partitions are misaligned
- assignment = asList("t1-0, t1-3, t2-0, t2-3, t2-6", "t1-1, t1-4, t2-1, t2-4, t3-0", "t1-2, t1-5, t2-2, t2-5, t3-1");
- List<String> owned = asList("t1-0, t1-1, t1-2, t1-3, t1-4", "t1-5, t2-0, t2-1, t2-2, t2-3", "t2-4, t2-5, t2-6, t3-0, t3-1");
- verifyRackAssignment(assignor, topics, 1, racks(3), racks(3), consumerTopics, owned, assignment, 0);
+ if (assignor instanceof CooperativeStickyAssignor) {
+ // Verify that rack-awareness is improved if already owned partitions are misaligned
+ assignment = asList("t1-0, t1-3", "t2-1", "t2-5, t3-1");
+ List<String> owned = asList("t1-0, t1-1, t1-2, t1-3, t1-4", "t1-5, t2-0, t2-1, t2-2, t2-3", "t2-4, t2-5, t2-6, t3-0, t3-1");
+ verifyRackAssignment(assignor, topics, 1, racks(3), racks(3), consumerTopics, owned, assignment, 0);
- // Verify that stickiness is retained when racks match
- AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
- verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, assignment, assignment, 0);
+ // Verify that stickiness is retained when racks match
+ assignment = asList("t1-0, t1-3, t2-0, t2-3, t2-6", "t1-1, t1-4, t2-1, t2-4, t3-0", "t1-2, t1-5, t2-2, t2-5, t3-1");
+ AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
+ verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, assignment, assignment, 0);
+ } else {
+ // Verify that rack-awareness is improved if already owned partitions are misaligned
+ assignment = asList("t1-0, t1-3, t2-0, t2-3, t2-6", "t1-1, t1-4, t2-1, t2-4, t3-0", "t1-2, t1-5, t2-2, t2-5, t3-1");
+ List<String> owned = asList("t1-0, t1-1, t1-2, t1-3, t1-4", "t1-5, t2-0, t2-1, t2-2, t2-3", "t2-4, t2-5, t2-6, t3-0, t3-1");
+ verifyRackAssignment(assignor, topics, 1, racks(3), racks(3), consumerTopics, owned, assignment, 0);
+
+ // Verify that stickiness is retained when racks match
+ AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
+ verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, assignment, assignment, 0);
+ }
}
private void verifyUniformSubscription(AbstractStickyAssignor assignor,