From e8d36070e0b4f799bd10b6388fa1ef85e8d51328 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 23 Mar 2023 13:02:31 -0700 Subject: [PATCH 01/44] Server Side Sticky Range Assignor full implementation --- checkstyle/suppressions.xml | 8 + .../group/assignor/AssignmentMemberSpec.java | 29 +- .../group/assignor/GroupAssignment.java | 4 + .../group/assignor/MemberAssignment.java | 33 +- .../ServerSideStickyRangeAssignor.java | 268 +++++++++++ .../ServerSideStickyRangeAssignorTest.java | 428 ++++++++++++++++++ 6 files changed, 741 insertions(+), 29 deletions(-) create mode 100644 group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index 93ae7c30d0f3e..3ff095f3760c9 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -320,6 +320,14 @@ + + + + + diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java index 2db33ddc74158..fe3d7c967e537 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java @@ -16,11 +16,13 @@ */ package org.apache.kafka.coordinator.group.assignor; -import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; -import java.util.Collection; -import java.util.Objects; +import java.util.List; +import java.util.Map; +import java.util.Set; import java.util.Optional; +import java.util.Objects; /** * The assignment specification for a consumer group member. @@ -37,29 +39,28 @@ public class AssignmentMemberSpec { final Optional rackId; /** - * The topics that the member is subscribed to. + * The topicIds of topics that the member is subscribed to. */ - final Collection subscribedTopics; + final List subscribedTopics; /** - * The current target partitions of the member. + * Maps the partitions assigned for this member per topicId */ - final Collection targetPartitions; + final Map> currentAssignmentPerTopic; public AssignmentMemberSpec( Optional instanceId, Optional rackId, - Collection subscribedTopics, - Collection targetPartitions + List subscribedTopics, + Map> currentAssignmentPerTopic ) { Objects.requireNonNull(instanceId); Objects.requireNonNull(rackId); Objects.requireNonNull(subscribedTopics); - Objects.requireNonNull(targetPartitions); this.instanceId = instanceId; this.rackId = rackId; this.subscribedTopics = subscribedTopics; - this.targetPartitions = targetPartitions; + this.currentAssignmentPerTopic = currentAssignmentPerTopic; } @Override @@ -72,7 +73,7 @@ public boolean equals(Object o) { if (!instanceId.equals(that.instanceId)) return false; if (!rackId.equals(that.rackId)) return false; if (!subscribedTopics.equals(that.subscribedTopics)) return false; - return targetPartitions.equals(that.targetPartitions); + return currentAssignmentPerTopic.equals(that.currentAssignmentPerTopic); } @Override @@ -80,7 +81,7 @@ public int hashCode() { int result = instanceId.hashCode(); result = 31 * result + rackId.hashCode(); result = 31 * result + subscribedTopics.hashCode(); - result = 31 * result + targetPartitions.hashCode(); + result = 31 * result + currentAssignmentPerTopic.hashCode(); return result; } @@ -89,7 +90,7 @@ public String toString() { return "AssignmentMemberSpec(instanceId=" + instanceId + ", rackId=" + rackId + ", subscribedTopics=" + subscribedTopics + - ", targetPartitions=" + targetPartitions + + ", currentAssignment=" + currentAssignmentPerTopic + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java index 5c0199aaec5a5..6256b7453a835 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java @@ -35,6 +35,10 @@ public GroupAssignment( this.members = members; } + public Map getMembers() { + return this.members; + } + @Override public boolean equals(Object o) { if (this == o) return true; diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java index 86c235d718771..3f9dd7f8d4378 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java @@ -16,25 +16,28 @@ */ package org.apache.kafka.coordinator.group.assignor; -import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; -import java.util.Collection; +import java.util.Map; import java.util.Objects; +import java.util.Set; + /** * The partition assignment for a consumer group member. */ public class MemberAssignment { - /** - * The target partitions assigned to this member. - */ - final Collection targetPartitions; - - public MemberAssignment( - Collection targetPartitions - ) { - Objects.requireNonNull(targetPartitions); - this.targetPartitions = targetPartitions; + + private final Map> assignmentPerTopic; + + + public MemberAssignment(Map> assignmentPerTopic) { + Objects.requireNonNull(assignmentPerTopic); + this.assignmentPerTopic = assignmentPerTopic; + } + + public Map> getAssignmentPerTopic() { + return this.assignmentPerTopic; } @Override @@ -44,16 +47,16 @@ public boolean equals(Object o) { MemberAssignment that = (MemberAssignment) o; - return targetPartitions.equals(that.targetPartitions); + return assignmentPerTopic.equals(that.assignmentPerTopic); } @Override public int hashCode() { - return targetPartitions.hashCode(); + return assignmentPerTopic.hashCode(); } @Override public String toString() { - return "MemberAssignment(targetPartitions=" + targetPartitions + ')'; + return "MemberAssignment(targetPartitions=" + assignmentPerTopic + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java new file mode 100644 index 0000000000000..186f8af1f0220 --- /dev/null +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java @@ -0,0 +1,268 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + *

The Server Side Sticky Range Assignor inherits properties of both the range assignor and the sticky assignor. + * Properties are :- + *

    + *
  • 1) Each consumer must get at least one partition per topic that it is subscribed to whenever the number of consumers is + * less than or equal to the number of partitions for that topic. (Range)
    • + *
  • 2) Partitions should be assigned to consumers in a way that facilitates join operations where required. (Range)
    • + * This can only be done if the topics are co-partitioned in the first place + * Co-partitioned:- + * Two streams are co-partitioned if the following conditions are met:- + * ->The keys must have the same schemas + * ->The topics involved must have the same number of partitions + *
  • 3) Consumers should retain as much as their previous assignment as possible. (Sticky)
    • + *
+ *

+ * + *

The algorithm works mainly in 5 steps described below + *

    + *
  • 1) Get a map of the consumersPerTopic created using the member subscriptions.
    • + *
  • 2) Get a list of consumers (potentiallyUnfilled) that have not met the minimum required quota for assignment AND + * get a list of sticky partitions that we want to retain in the new assignment.
    • + *
  • 3) Add consumers from potentiallyUnfilled to Unfilled if they haven't met the total required quota = minQuota + (if necessary) extraPartition
    • + *
  • 4) Get a list of available partitions by calculating the difference between total partitions and assigned sticky partitions
    • + *
  • 5) Iterate through unfilled consumers and assign partitions from available partitions
    • + *
+ *

+ * + */ +package org.apache.kafka.coordinator.group.assignor; + + +import org.apache.kafka.common.Uuid; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.HashMap; +import java.util.Collections; + +import static java.lang.Math.min; + +public class ServerSideStickyRangeAssignor implements PartitionAssignor { + + public static final String RANGE_ASSIGNOR_NAME = "range-sticky"; + + @Override + public String name() { + return RANGE_ASSIGNOR_NAME; + } + + protected static void putList(Map> map, K key, V value) { + List list = map.computeIfAbsent(key, k -> new ArrayList<>()); + list.add(value); + } + + protected static void putSet(Map> map, K key, V value) { + Set set = map.computeIfAbsent(key, k -> new HashSet<>()); + set.add(value); + } + + static class Pair { + private final T first; + private final U second; + + public Pair(T first, U second) { + this.first = first; + this.second = second; + } + + public T getFirst() { + return first; + } + + public U getSecond() { + return second; + } + + @Override + public String toString() { + return "(" + first + ", " + second + ")"; + } + } + + // Returns a map of the list of consumers per Topic (keyed by topicId) + private Map> consumersPerTopic(AssignmentSpec assignmentSpec) { + Map> mapTopicsToConsumers = new HashMap<>(); + Map membersData = assignmentSpec.members; + + for (Map.Entry memberEntry : membersData.entrySet()) { + String memberId = memberEntry.getKey(); + AssignmentMemberSpec memberMetadata = memberEntry.getValue(); + List topics = new ArrayList<>(memberMetadata.subscribedTopics); + for (Uuid topicId: topics) { + putList(mapTopicsToConsumers, topicId, memberId); + } + } + return mapTopicsToConsumers; + } + + private Map> getAvailablePartitionsPerTopic(AssignmentSpec assignmentSpec, Map> assignedStickyPartitionsPerTopic) { + Map> availablePartitionsPerTopic = new HashMap<>(); + Map topicsMetadata = assignmentSpec.topics; + // Iterate through the topics map provided in assignmentSpec + for (Map.Entry topicMetadataEntry : topicsMetadata.entrySet()) { + Uuid topicId = topicMetadataEntry.getKey(); + availablePartitionsPerTopic.put(topicId, new ArrayList<>()); + int numPartitions = topicsMetadata.get(topicId).numPartitions; + // since the loop iterates from 0 to n, the partitions will be in ascending order within the list of available partitions per topic + Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.get(topicId); + for (int i = 0; i < numPartitions; i++) { + if (assignedStickyPartitionsForTopic == null || !assignedStickyPartitionsForTopic.contains(i)) { + availablePartitionsPerTopic.get(topicId).add(i); + } + } + } + return availablePartitionsPerTopic; + } + + @Override + public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAssignorException { + Map>> membersWithNewAssignmentPerTopic = new HashMap<>(); + // Step 1 + Map> consumersPerTopic = consumersPerTopic(assignmentSpec); + // Step 2 + Map>> unfilledConsumersPerTopic = new HashMap<>(); + Map> assignedStickyPartitionsPerTopic = new HashMap<>(); + + for (Map.Entry> topicEntry : consumersPerTopic.entrySet()) { + Uuid topicId = topicEntry.getKey(); + // For each topic we have a temporary list of consumers stored in potentiallyUnfilledConsumers. + // The list is populated with consumers that satisfy one of the two conditions :- + // 1) Consumers that have the minimum required number of partitions .i.e numPartitionsPerConsumer BUT they could be assigned an extra partition later on. + // In this case we add the consumer to the unfilled consumers map iff an extra partition needs to be assigned to it. + // 2) Consumers that don't have the minimum required partitions, so irrespective of whether they get an extra partition or not they get added to the unfilled map later. + List> potentiallyUnfilledConsumers = new ArrayList<>(); + List consumersForTopic = topicEntry.getValue(); + + AssignmentTopicMetadata topicData = assignmentSpec.topics.get(topicId); + int numPartitionsForTopic = topicData.numPartitions; + int numPartitionsPerConsumer = numPartitionsForTopic / consumersForTopic.size(); + // Each consumer can get only one extra partition per topic after receiving the minimum quota = numPartitionsPerConsumer + int numConsumersWithExtraPartition = numPartitionsForTopic % consumersForTopic.size(); + + for (String memberId: consumersForTopic) { + // Size of the older assignment, this will be 0 when assign is called for the first time. + // The older assignment is required when a reassignment occurs to ensure stickiness. + int currentAssignmentSize = 0; + List currentAssignmentListForTopic = new ArrayList<>(); + // Convert the set to a list first and sort the partitions in numeric order since we want the same partition numbers from each topic + // to go to the same consumer in case of co-partitioned topics. + Set currentAssignmentSetForTopic = assignmentSpec.members.get(memberId).currentAssignmentPerTopic.get(topicId); + // We want to make sure that the currentAssignment is not null to avoid a null pointer exception + if (currentAssignmentSetForTopic != null) { + currentAssignmentListForTopic = new ArrayList<>(currentAssignmentSetForTopic); + currentAssignmentSize = currentAssignmentListForTopic.size(); + } + // Initially, minRequiredQuota is the minimum number of partitions that each consumer should get. + // The value is at least 1 unless numConsumers subscribed is greater than numPartitions. In such cases, all consumers get assigned "extra partitions". + int minRequiredQuota = numPartitionsPerConsumer; + // We need to make sure that the partitions we're keeping are still part of the topic metadata. + // Ex:- In case the number of partitions is reduced for a topic, the sticky partitions must be part of the new set of partitions available. + for (int i = 0; i < currentAssignmentListForTopic.size() && currentAssignmentSize > 0; i++) { + // numPartitionsForTopic - 1 is the max possible partition number + if (currentAssignmentListForTopic.get(i) >= numPartitionsForTopic) { + currentAssignmentSize--; + } + } + // If there are previously assigned partitions present, we want to retain + if (currentAssignmentSize > 0) { + // We either need to retain currentSize number of partitions when currentSize < required OR required number of partitions otherwise. + int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); + for (int i = 0; i < retainedPartitionsCount; i++) { + Collections.sort(currentAssignmentListForTopic); + putSet(assignedStickyPartitionsPerTopic, topicId, currentAssignmentListForTopic.get(i)); + membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); + } + } + // Number of partitions left to reach the minRequiredQuota + int remaining = minRequiredQuota - currentAssignmentSize; + + // There are 3 cases w.r.t value of remaining + // 1) remaining < 0 this means that the consumer has more than the min required amount. + // It could have an extra partition, so we check for that. + if (remaining < 0 && numConsumersWithExtraPartition > 0) { + // In order to remain as sticky as possible, since the order of members can be different, we want the consumers that already had extra + // partitions to retain them if it's still required, instead of assigning the extras to the first few consumers directly. + // Ex:- If two consumers out of 3 are supposed to get an extra partition and currently 1 of them already has the extra, we want this consumer + // to retain it first and later if we have partitions left they will be assigned to the first few consumers and the unfilled map is updated + numConsumersWithExtraPartition--; + // Since we already added the minimumRequiredQuota of partitions in the previous step (until minReq - 1), we just need to + // add the extra partition which will be present at the index right after min quota is satisfied. + putSet(assignedStickyPartitionsPerTopic, topicId, currentAssignmentListForTopic.get(minRequiredQuota)); + membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(minRequiredQuota)); + } else { + // 3) If remaining = 0 it has min req partitions but there is scope for getting an extra partition later on, so it is a potentialUnfilledConsumer. + // 4) If remaining > 0 it doesn't even have the min required partitions, and it definitely is unfilled so it should be added to potentialUnfilledConsumers. + Pair newPair = new Pair<>(memberId, remaining); + potentiallyUnfilledConsumers.add(newPair); + } + } + + // Step 3 + // Iterate through potentially unfilled consumers and if remaining > 0 after considering the extra partitions assignment, add to the unfilled list. + for (Pair pair : potentiallyUnfilledConsumers) { + String memberId = pair.getFirst(); + Integer remaining = pair.getSecond(); + if (numConsumersWithExtraPartition > 0) { + remaining++; + numConsumersWithExtraPartition--; + } + if (remaining > 0) { + Pair newPair = new Pair<>(memberId, remaining); + putList(unfilledConsumersPerTopic, topicId, newPair); + } + } + } + + // Step 4 + // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get available partitions. + Map> availablePartitionsPerTopic = getAvailablePartitionsPerTopic(assignmentSpec, assignedStickyPartitionsPerTopic); + + // Step 5 + // Iterate through the unfilled consumers list and assign remaining number of partitions from the availablePartitions list + for (Map.Entry>> unfilledEntry : unfilledConsumersPerTopic.entrySet()) { + Uuid topicId = unfilledEntry.getKey(); + List> unfilledConsumersForTopic = unfilledEntry.getValue(); + + for (Pair stringIntegerPair : unfilledConsumersForTopic) { + String memberId = stringIntegerPair.getFirst(); + int remaining = stringIntegerPair.getSecond(); + // assign the first few partitions from the list and then delete them from the availablePartitions list for this topic + List subList = availablePartitionsPerTopic.get(topicId).subList(0, remaining); + membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).addAll(subList); + availablePartitionsPerTopic.get(topicId).removeAll(subList); + } + } + + // Consolidate the maps into MemberAssignment and then finally map each consumer to a MemberAssignment. + Map membersWithNewAssignment = new HashMap<>(); + for (Map.Entry>> consumer : membersWithNewAssignmentPerTopic.entrySet()) { + String consumerId = consumer.getKey(); + Map> assignmentPerTopic = consumer.getValue(); + membersWithNewAssignment.computeIfAbsent(consumerId, k -> new MemberAssignment(assignmentPerTopic)); + } + + return new GroupAssignment(membersWithNewAssignment); + } +} + diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java new file mode 100644 index 0000000000000..228157280f610 --- /dev/null +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java @@ -0,0 +1,428 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.coordinator.group.assignor; + + +import org.apache.kafka.common.Uuid; + +import org.junit.jupiter.api.Test; + +import java.util.Map; +import java.util.Set; + + + +import java.util.*; + +import static org.junit.jupiter.api.Assertions.*; + +public class ServerSideStickyRangeAssignorTest { + + private final ServerSideStickyRangeAssignor assignor = new ServerSideStickyRangeAssignor(); + + private final String topic1Name = "topic1"; + private final Uuid topic1Uuid = Uuid.randomUuid(); + + private final String topic2Name = "topic2"; + private final Uuid topic2Uuid = Uuid.randomUuid(); + + private final String topic3Name = "topic3"; + private final Uuid topic3Uuid = Uuid.randomUuid(); + + private final String consumerA = "A"; + private final String consumerB = "B"; + private final String consumerC = "C"; + + @Test + public void testOneConsumerNoTopic() { + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + Map members = new HashMap<>(); + List subscribedTopics = new ArrayList<>(); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopics, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment groupAssignment = assignor.assign(assignmentSpec); + + assertTrue(groupAssignment.getMembers().isEmpty()); + } + + @Test + public void testFirstAssignmentTwoConsumersTwoTopicsSameSubscriptions() { + // A -> T1, T3 // B -> T1, T3 // T1 -> 3 Partitions // T3 -> 2 Partitions + // Topics + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + // Members + Map members = new HashMap<>(); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + // Topic 3 Partitions Assignment + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + assertAssignment(expectedAssignment, computedAssignment); + } + + @Test + public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() { + // A -> T1, T2 // B -> T3 // C -> T2, T3 // T1 -> 3 Partitions // T2 -> 3 Partitions // T3 -> 2 Partitions + // Topics + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + // Members + Map members = new HashMap<>(); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); + // Consumer B + List subscribedTopicsB = new ArrayList<>(Collections.singletonList(topic3Uuid)); + members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); + // Consumer B + List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic2Uuid, topic3Uuid)); + members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + // Topic 3 Partitions Assignment + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + assertAssignment(expectedAssignment, computedAssignment); + } + + @Test + public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { + // Topics + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + // Topic 3 has 2 partitions but three consumers subscribed to it - one of them will not get an assignment + topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + // Members + Map members = new HashMap<>(); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); + // Consumer C + List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + assertAssignment(expectedAssignment, computedAssignment); + } + + @Test + public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerAdded() { + // Topics + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 2)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 2)); + // Members + Map members = new HashMap<>(); + // Add a new consumer to trigger a re-assignment + // Consumer C + List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Collections.singletonList(0))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Collections.singletonList(0))); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(1))); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(1))); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + // Consumer C shouldn't get any assignment, due to stickiness A, B retain their assignments + assertNull(computedAssignment.getMembers().get(consumerC)); + assertAssignment(expectedAssignment, computedAssignment); + assertCoPartitionJoinProperty(computedAssignment); + } + + @Test + public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { + // T1, T2 both have 3 partitions each when first assignment was calculated -> currentAssignmentForX + Map members = new HashMap<>(); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + + // Simulating adding a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 4)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 4)); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(2, 3))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(2, 3))); + + assertAssignment(expectedAssignment, computedAssignment); + assertCoPartitionJoinProperty(computedAssignment); + } + + @Test + public void testReassignmentWhenOnePartitionRemovedForTwoConsumersTwoTopics() { + Map members = new HashMap<>(); + // T1, T2 both have 3 partitions each initially and after removal have 2 partitions each + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + + // Remove a partition from both topic1 and topic 2 -> total partitions = 2 + // Simulating removing a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 2)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 2)); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + assertAssignment(expectedAssignment, computedAssignment); + assertCoPartitionJoinProperty(computedAssignment); + } + + @Test + public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoConsumersTwoTopics() { + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); + + Map members = new HashMap<>(); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + // Add a new consumer to trigger a re-assignment + // Consumer C + List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + assertAssignment(expectedAssignment, computedAssignment); + assertCoPartitionJoinProperty(computedAssignment); + } + + @Test + public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoConsumersTwoTopics() { + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); + + Map members = new HashMap<>(); + // Consumer A was removed + + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); + + assertAssignment(expectedAssignment, computedAssignment); + assertCoPartitionJoinProperty(computedAssignment); + } + + @Test + public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignmentWithThreeConsumersTwoTopics() { + + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); + topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + + // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3 + // Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2 + Map members = new HashMap<>(); + // Consumer A + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1, 2))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Collections.singletonList(0))); + // Change subscriptions + List subscribedTopicsA = new ArrayList<>(Collections.singletonList(topic1Uuid)); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + // Consumer B + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(1))); + // Change subscriptions + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid, topic3Uuid)); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + // Consumer C + Map> currentAssignmentForC = new HashMap<>(); + currentAssignmentForC.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForC.put(topic3Uuid, new HashSet<>(Arrays.asList(0, 1))); + // Change subscriptions + List subscribedTopicsC = new ArrayList<>(Collections.singletonList(topic2Uuid)); + members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + // Topic 3 Partitions Assignment + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + + assertAssignment(expectedAssignment, computedAssignment); + } + + // We have a set of sets with the partitions that should be distributed amongst the consumers, if it exists then remove it from the set. + private void assertAssignment(Map>> expectedAssignment, GroupAssignment computedGroupAssignment) { + for (Map.Entry member : computedGroupAssignment.getMembers().entrySet()) { + Map> computedAssignmentForMember = member.getValue().getAssignmentPerTopic(); + for (Map.Entry> assignmentForTopic : computedAssignmentForMember.entrySet()) { + Uuid topicId = assignmentForTopic.getKey(); + Set assignmentPartitionsSet = assignmentForTopic.getValue(); + assertTrue(expectedAssignment.get(topicId).contains(assignmentPartitionsSet)); + expectedAssignment.remove(assignmentPartitionsSet); + } + } + } + + private void assertCoPartitionJoinProperty(GroupAssignment groupAssignment) { + for (Map.Entry member : groupAssignment.getMembers().entrySet()) { + Map> computedAssignmentForMember = member.getValue().getAssignmentPerTopic(); + Set compareSet = new HashSet<>(); + for (Map.Entry> topicAssignment : computedAssignmentForMember.entrySet()) { + Set partitionsForTopicSet = topicAssignment.getValue(); + if (compareSet.isEmpty()) { + compareSet = partitionsForTopicSet; + } + assertEquals(compareSet, partitionsForTopicSet); + } + } + } +} From 3df34605e078beb9babcbeadc7b119a70da63796 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Mon, 27 Mar 2023 20:17:14 -0700 Subject: [PATCH 02/44] removed reduce partitions case --- .../ServerSideStickyRangeAssignor.java | 19 +++------ .../ServerSideStickyRangeAssignorTest.java | 39 ------------------- 2 files changed, 6 insertions(+), 52 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java index 186f8af1f0220..e9d90ccfb7848 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java @@ -108,7 +108,7 @@ private Map> consumersPerTopic(AssignmentSpec assignmentSpec) for (Map.Entry memberEntry : membersData.entrySet()) { String memberId = memberEntry.getKey(); AssignmentMemberSpec memberMetadata = memberEntry.getValue(); - List topics = new ArrayList<>(memberMetadata.subscribedTopics); + List topics = memberMetadata.subscribedTopics; for (Uuid topicId: topics) { putList(mapTopicsToConsumers, topicId, memberId); } @@ -176,20 +176,13 @@ public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAss // Initially, minRequiredQuota is the minimum number of partitions that each consumer should get. // The value is at least 1 unless numConsumers subscribed is greater than numPartitions. In such cases, all consumers get assigned "extra partitions". int minRequiredQuota = numPartitionsPerConsumer; - // We need to make sure that the partitions we're keeping are still part of the topic metadata. - // Ex:- In case the number of partitions is reduced for a topic, the sticky partitions must be part of the new set of partitions available. - for (int i = 0; i < currentAssignmentListForTopic.size() && currentAssignmentSize > 0; i++) { - // numPartitionsForTopic - 1 is the max possible partition number - if (currentAssignmentListForTopic.get(i) >= numPartitionsForTopic) { - currentAssignmentSize--; - } - } + // If there are previously assigned partitions present, we want to retain if (currentAssignmentSize > 0) { // We either need to retain currentSize number of partitions when currentSize < required OR required number of partitions otherwise. int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); + Collections.sort(currentAssignmentListForTopic); for (int i = 0; i < retainedPartitionsCount; i++) { - Collections.sort(currentAssignmentListForTopic); putSet(assignedStickyPartitionsPerTopic, topicId, currentAssignmentListForTopic.get(i)); membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); } @@ -243,14 +236,14 @@ public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAss for (Map.Entry>> unfilledEntry : unfilledConsumersPerTopic.entrySet()) { Uuid topicId = unfilledEntry.getKey(); List> unfilledConsumersForTopic = unfilledEntry.getValue(); - + int newStartPointer = 0; for (Pair stringIntegerPair : unfilledConsumersForTopic) { String memberId = stringIntegerPair.getFirst(); int remaining = stringIntegerPair.getSecond(); // assign the first few partitions from the list and then delete them from the availablePartitions list for this topic - List subList = availablePartitionsPerTopic.get(topicId).subList(0, remaining); + List subList = availablePartitionsPerTopic.get(topicId).subList(newStartPointer, newStartPointer + remaining); + newStartPointer += remaining; membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).addAll(subList); - availablePartitionsPerTopic.get(topicId).removeAll(subList); } } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java index 228157280f610..54a73229718bb 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java @@ -242,45 +242,6 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { assertCoPartitionJoinProperty(computedAssignment); } - @Test - public void testReassignmentWhenOnePartitionRemovedForTwoConsumersTwoTopics() { - Map members = new HashMap<>(); - // T1, T2 both have 3 partitions each initially and after removal have 2 partitions each - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); - - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); - - // Remove a partition from both topic1 and topic 2 -> total partitions = 2 - // Simulating removing a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 2)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 2)); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - - Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - - assertAssignment(expectedAssignment, computedAssignment); - assertCoPartitionJoinProperty(computedAssignment); - } - @Test public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoConsumersTwoTopics() { Map topics = new HashMap<>(); From 74105291cc687fe41b75de692fc624ab1335341a Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 28 Mar 2023 12:21:33 -0700 Subject: [PATCH 03/44] removed reduce partitions case --- .../group/assignor/ServerSideStickyRangeAssignor.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java index e9d90ccfb7848..8c5f021405bde 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java @@ -60,7 +60,7 @@ public class ServerSideStickyRangeAssignor implements PartitionAssignor { - public static final String RANGE_ASSIGNOR_NAME = "range-sticky"; + public static final String RANGE_ASSIGNOR_NAME = "range"; @Override public String name() { From 3b6e65ea3b8452dd923991c3dba49d370c8f0d4e Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Wed, 29 Mar 2023 12:30:13 -0700 Subject: [PATCH 04/44] Addressed PR comments --- .../ServerSideStickyRangeAssignor.java | 72 +++++++++---------- 1 file changed, 34 insertions(+), 38 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java index 8c5f021405bde..c2e2dcae54c5c 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java @@ -14,6 +14,20 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +package org.apache.kafka.coordinator.group.assignor; + + +import org.apache.kafka.common.Uuid; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.HashMap; +import java.util.Collections; + +import static java.lang.Math.min; /** *

The Server Side Sticky Range Assignor inherits properties of both the range assignor and the sticky assignor. @@ -43,21 +57,6 @@ *

* */ -package org.apache.kafka.coordinator.group.assignor; - - -import org.apache.kafka.common.Uuid; - -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.HashMap; -import java.util.Collections; - -import static java.lang.Math.min; - public class ServerSideStickyRangeAssignor implements PartitionAssignor { public static final String RANGE_ASSIGNOR_NAME = "range"; @@ -67,12 +66,12 @@ public String name() { return RANGE_ASSIGNOR_NAME; } - protected static void putList(Map> map, K key, V value) { + private static void putList(Map> map, K key, V value) { List list = map.computeIfAbsent(key, k -> new ArrayList<>()); list.add(value); } - protected static void putSet(Map> map, K key, V value) { + private static void putSet(Map> map, K key, V value) { Set set = map.computeIfAbsent(key, k -> new HashSet<>()); set.add(value); } @@ -119,18 +118,19 @@ private Map> consumersPerTopic(AssignmentSpec assignmentSpec) private Map> getAvailablePartitionsPerTopic(AssignmentSpec assignmentSpec, Map> assignedStickyPartitionsPerTopic) { Map> availablePartitionsPerTopic = new HashMap<>(); Map topicsMetadata = assignmentSpec.topics; - // Iterate through the topics map provided in assignmentSpec + for (Map.Entry topicMetadataEntry : topicsMetadata.entrySet()) { Uuid topicId = topicMetadataEntry.getKey(); - availablePartitionsPerTopic.put(topicId, new ArrayList<>()); + ArrayList availablePartitionsForTopic = new ArrayList<>(); int numPartitions = topicsMetadata.get(topicId).numPartitions; // since the loop iterates from 0 to n, the partitions will be in ascending order within the list of available partitions per topic - Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.get(topicId); + Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.getOrDefault(topicId, new HashSet<>()); for (int i = 0; i < numPartitions; i++) { - if (assignedStickyPartitionsForTopic == null || !assignedStickyPartitionsForTopic.contains(i)) { - availablePartitionsPerTopic.get(topicId).add(i); + if (!assignedStickyPartitionsForTopic.contains(i)) { + availablePartitionsForTopic.add(i); } } + availablePartitionsPerTopic.put(topicId, availablePartitionsForTopic); } return availablePartitionsPerTopic; } @@ -156,28 +156,23 @@ public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAss AssignmentTopicMetadata topicData = assignmentSpec.topics.get(topicId); int numPartitionsForTopic = topicData.numPartitions; - int numPartitionsPerConsumer = numPartitionsForTopic / consumersForTopic.size(); + // Initially, minRequiredQuota is the minimum number of partitions that each consumer should get i.e numPartitionsPerConsumer. + // Idle consumers case :- The numConsumers subscribed to a topic is greater than numPartitions. In such cases, all consumers get assigned via the "extra partitions" logic since min = 0. + int minRequiredQuota = numPartitionsForTopic / consumersForTopic.size(); // Each consumer can get only one extra partition per topic after receiving the minimum quota = numPartitionsPerConsumer int numConsumersWithExtraPartition = numPartitionsForTopic % consumersForTopic.size(); for (String memberId: consumersForTopic) { - // Size of the older assignment, this will be 0 when assign is called for the first time. - // The older assignment is required when a reassignment occurs to ensure stickiness. - int currentAssignmentSize = 0; - List currentAssignmentListForTopic = new ArrayList<>(); + // Convert the set to a list first and sort the partitions in numeric order since we want the same partition numbers from each topic // to go to the same consumer in case of co-partitioned topics. - Set currentAssignmentSetForTopic = assignmentSpec.members.get(memberId).currentAssignmentPerTopic.get(topicId); - // We want to make sure that the currentAssignment is not null to avoid a null pointer exception - if (currentAssignmentSetForTopic != null) { - currentAssignmentListForTopic = new ArrayList<>(currentAssignmentSetForTopic); - currentAssignmentSize = currentAssignmentListForTopic.size(); - } - // Initially, minRequiredQuota is the minimum number of partitions that each consumer should get. - // The value is at least 1 unless numConsumers subscribed is greater than numPartitions. In such cases, all consumers get assigned "extra partitions". - int minRequiredQuota = numPartitionsPerConsumer; + Set currentAssignmentSetForTopic = assignmentSpec.members.get(memberId).currentAssignmentPerTopic.getOrDefault(topicId, new HashSet<>()); + // Size of the older assignment, this will be 0 when assign is called for the first time. + // The older assignment is required when a reassignment occurs to ensure stickiness. + int currentAssignmentSize = currentAssignmentSetForTopic.size(); + List currentAssignmentListForTopic = new ArrayList<>(currentAssignmentSetForTopic); - // If there are previously assigned partitions present, we want to retain + // If there are previously assigned partitions present, we want to retain them. if (currentAssignmentSize > 0) { // We either need to retain currentSize number of partitions when currentSize < required OR required number of partitions otherwise. int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); @@ -187,7 +182,8 @@ public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAss membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); } } - // Number of partitions left to reach the minRequiredQuota + + // Number of partitions left to reach the minRequiredQuota. int remaining = minRequiredQuota - currentAssignmentSize; // There are 3 cases w.r.t value of remaining From 4b321aa5374427a15346be9b87b7c545ccb7db61 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Mon, 3 Apr 2023 11:54:59 -0700 Subject: [PATCH 05/44] Addressed PR comments --- .../assignor/ServerSideStickyRangeAssignorTest.java | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java index 54a73229718bb..5704d68cc5fde 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java @@ -108,7 +108,7 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() // Consumer B List subscribedTopicsB = new ArrayList<>(Collections.singletonList(topic3Uuid)); members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); - // Consumer B + // Consumer C List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic2Uuid, topic3Uuid)); members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); @@ -362,8 +362,8 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme // We have a set of sets with the partitions that should be distributed amongst the consumers, if it exists then remove it from the set. private void assertAssignment(Map>> expectedAssignment, GroupAssignment computedGroupAssignment) { - for (Map.Entry member : computedGroupAssignment.getMembers().entrySet()) { - Map> computedAssignmentForMember = member.getValue().getAssignmentPerTopic(); + for (MemberAssignment member : computedGroupAssignment.getMembers().values()) { + Map> computedAssignmentForMember = member.getAssignmentPerTopic(); for (Map.Entry> assignmentForTopic : computedAssignmentForMember.entrySet()) { Uuid topicId = assignmentForTopic.getKey(); Set assignmentPartitionsSet = assignmentForTopic.getValue(); @@ -374,11 +374,10 @@ private void assertAssignment(Map>> expectedAssignment, G } private void assertCoPartitionJoinProperty(GroupAssignment groupAssignment) { - for (Map.Entry member : groupAssignment.getMembers().entrySet()) { - Map> computedAssignmentForMember = member.getValue().getAssignmentPerTopic(); + for (MemberAssignment member : groupAssignment.getMembers().values()) { + Map> computedAssignmentForMember = member.getAssignmentPerTopic(); Set compareSet = new HashSet<>(); - for (Map.Entry> topicAssignment : computedAssignmentForMember.entrySet()) { - Set partitionsForTopicSet = topicAssignment.getValue(); + for (Set partitionsForTopicSet : computedAssignmentForMember.values()) { if (compareSet.isEmpty()) { compareSet = partitionsForTopicSet; } From 7352cbb5683f676382f56e6ae5c926f3cae98824 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 6 Apr 2023 10:55:06 -0700 Subject: [PATCH 06/44] minor --- .../group/assignor/ServerSideStickyRangeAssignorTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java index 5704d68cc5fde..73a7860239e0a 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java @@ -98,7 +98,7 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() // Topics Map topics = new HashMap<>(); topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); // Members Map members = new HashMap<>(); From 57b94ca208cb374b22d61eb044ef8004ea09479f Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 6 Apr 2023 11:12:16 -0700 Subject: [PATCH 07/44] Made subscribed topics a collection, removed * import exception --- checkstyle/suppressions.xml | 2 -- .../group/assignor/AssignmentMemberSpec.java | 7 ++++--- .../assignor/ServerSideStickyRangeAssignor.java | 12 ++++++------ .../ServerSideStickyRangeAssignorTest.java | 16 ++++++++++------ 4 files changed, 20 insertions(+), 17 deletions(-) diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index 1bf330e52fc73..e2a501ae3a575 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -323,8 +323,6 @@ files="ServerSideStickyRangeAssignor.java"/> - subscribedTopics; + final Collection subscribedTopics; /** * Maps the partitions assigned for this member per topicId diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java index c2e2dcae54c5c..bb89296d0c0ee 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java @@ -16,16 +16,16 @@ */ package org.apache.kafka.coordinator.group.assignor; - import org.apache.kafka.common.Uuid; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.HashMap; -import java.util.Collections; import static java.lang.Math.min; @@ -107,7 +107,7 @@ private Map> consumersPerTopic(AssignmentSpec assignmentSpec) for (Map.Entry memberEntry : membersData.entrySet()) { String memberId = memberEntry.getKey(); AssignmentMemberSpec memberMetadata = memberEntry.getValue(); - List topics = memberMetadata.subscribedTopics; + Collection topics = memberMetadata.subscribedTopics; for (Uuid topicId: topics) { putList(mapTopicsToConsumers, topicId, memberId); } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java index 73a7860239e0a..52d36eece5610 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java @@ -17,19 +17,23 @@ package org.apache.kafka.coordinator.group.assignor; - import org.apache.kafka.common.Uuid; import org.junit.jupiter.api.Test; +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.Set; - - -import java.util.*; - -import static org.junit.jupiter.api.Assertions.*; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; public class ServerSideStickyRangeAssignorTest { From 5065109332845541c5591a52683f21d37982d80a Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 6 Apr 2023 13:43:26 -0700 Subject: [PATCH 08/44] Interface changes for assignment map and subscribed topics, added new TopicIdToPartition data structure --- .../group/assignor/AssignmentMemberSpec.java | 27 +++++----- .../group/assignor/MemberAssignment.java | 27 +++++----- .../group/common/TopicIdToPartition.java | 50 +++++++++++++++++++ 3 files changed, 80 insertions(+), 24 deletions(-) create mode 100644 group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java index 2db33ddc74158..b3e63994d980b 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java @@ -16,11 +16,13 @@ */ package org.apache.kafka.coordinator.group.assignor; -import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; import java.util.Collection; +import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.Set; /** * The assignment specification for a consumer group member. @@ -37,29 +39,30 @@ public class AssignmentMemberSpec { final Optional rackId; /** - * The topics that the member is subscribed to. + * The topicIds of topics that the member is subscribed to. */ - final Collection subscribedTopics; + final Collection subscribedTopics; /** - * The current target partitions of the member. + * Partitions assigned for this member grouped by topicId */ - final Collection targetPartitions; + final Map> currentAssignmentTopicIdPartitions; public AssignmentMemberSpec( + Optional instanceId, Optional rackId, - Collection subscribedTopics, - Collection targetPartitions + Collection subscribedTopics, + Map> currentAssignmentTopicIdPartitions ) { Objects.requireNonNull(instanceId); Objects.requireNonNull(rackId); Objects.requireNonNull(subscribedTopics); - Objects.requireNonNull(targetPartitions); + Objects.requireNonNull(currentAssignmentTopicIdPartitions); this.instanceId = instanceId; this.rackId = rackId; this.subscribedTopics = subscribedTopics; - this.targetPartitions = targetPartitions; + this.currentAssignmentTopicIdPartitions = currentAssignmentTopicIdPartitions; } @Override @@ -72,7 +75,7 @@ public boolean equals(Object o) { if (!instanceId.equals(that.instanceId)) return false; if (!rackId.equals(that.rackId)) return false; if (!subscribedTopics.equals(that.subscribedTopics)) return false; - return targetPartitions.equals(that.targetPartitions); + return currentAssignmentTopicIdPartitions.equals(that.currentAssignmentTopicIdPartitions); } @Override @@ -80,7 +83,7 @@ public int hashCode() { int result = instanceId.hashCode(); result = 31 * result + rackId.hashCode(); result = 31 * result + subscribedTopics.hashCode(); - result = 31 * result + targetPartitions.hashCode(); + result = 31 * result + currentAssignmentTopicIdPartitions.hashCode(); return result; } @@ -89,7 +92,7 @@ public String toString() { return "AssignmentMemberSpec(instanceId=" + instanceId + ", rackId=" + rackId + ", subscribedTopics=" + subscribedTopics + - ", targetPartitions=" + targetPartitions + + ", targetPartitions=" + currentAssignmentTopicIdPartitions + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java index 86c235d718771..fa05e4398222b 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java @@ -16,25 +16,28 @@ */ package org.apache.kafka.coordinator.group.assignor; -import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; -import java.util.Collection; +import java.util.Map; import java.util.Objects; +import java.util.Set; /** * The partition assignment for a consumer group member. */ public class MemberAssignment { /** - * The target partitions assigned to this member. + * The target partitions assigned to this member grouped by topicId. */ - final Collection targetPartitions; + private final Map> topicIdPartitionsMap; - public MemberAssignment( - Collection targetPartitions - ) { - Objects.requireNonNull(targetPartitions); - this.targetPartitions = targetPartitions; + public MemberAssignment(Map> topicIdPartitionsForAssignment) { + Objects.requireNonNull(topicIdPartitionsForAssignment); + this.topicIdPartitionsMap = topicIdPartitionsForAssignment; + } + + public Map> getTopicIdPartitionsMap() { + return this.topicIdPartitionsMap; } @Override @@ -44,16 +47,16 @@ public boolean equals(Object o) { MemberAssignment that = (MemberAssignment) o; - return targetPartitions.equals(that.targetPartitions); + return topicIdPartitionsMap.equals(that.topicIdPartitionsMap); } @Override public int hashCode() { - return targetPartitions.hashCode(); + return topicIdPartitionsMap.hashCode(); } @Override public String toString() { - return "MemberAssignment(targetPartitions=" + targetPartitions + ')'; + return "MemberAssignment ( Assignment per topic Id = " + topicIdPartitionsMap + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java new file mode 100644 index 0000000000000..c84e0c1384123 --- /dev/null +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.coordinator.group.common; + +import org.apache.kafka.common.Uuid; + +import java.util.List; +import java.util.Objects; +import java.util.Optional; + +public class TopicIdToPartition { + private final Uuid topicId; + private final Integer partition; + private final Optional> rackIds; + + public TopicIdToPartition(Uuid topicId, Integer topicPartition, Optional> rackIds) { + this.topicId = Objects.requireNonNull(topicId, "topicId can not be null"); + this.partition = Objects.requireNonNull(topicPartition, "topicPartition can not be null"); + this.rackIds = rackIds; + } + + /** + * @return Universally unique id representing this topic partition. + */ + public Uuid topicId() { + return topicId; + } + + /** + * @return the partition number. + */ + public int partition() { + return partition; + } + +} From 100a04e5e064be0bda392ca8f02cb7a7e89f16c8 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Fri, 7 Apr 2023 11:25:33 -0700 Subject: [PATCH 09/44] Added toString method and hash --- .../group/common/TopicIdToPartition.java | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java index c84e0c1384123..ea7d1114feb32 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java @@ -47,4 +47,30 @@ public int partition() { return partition; } + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + TopicIdToPartition that = (TopicIdToPartition) o; + return topicId.equals(that.topicId) && + partition.equals(that.partition); + } + + @Override + public int hashCode() { + final int prime = 31; + int result = prime + topicId.hashCode(); + result = prime * result + partition.hashCode(); + return result; + } + + @Override + public String toString() { + return topicId() + "-" + partition(); + } + } From 73c7fdcaf824a31af442b1a4aa2db81b37cc5a9e Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 11 Apr 2023 15:09:44 -0700 Subject: [PATCH 10/44] Made all attributes private and added getter methods, changed names according to PR comments --- .../group/assignor/AssignmentMemberSpec.java | 39 +++++++++++++------ .../group/assignor/AssignmentSpec.java | 12 +++++- .../assignor/AssignmentTopicMetadata.java | 12 +++++- .../group/assignor/GroupAssignment.java | 6 ++- .../group/assignor/MemberAssignment.java | 16 ++++---- .../group/common/TopicIdToPartition.java | 6 +-- 6 files changed, 63 insertions(+), 28 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java index b3e63994d980b..33dfe71165c85 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java @@ -31,38 +31,53 @@ public class AssignmentMemberSpec { /** * The instance ID if provided. */ - final Optional instanceId; + private final Optional instanceId; /** * The rack ID if provided. */ - final Optional rackId; + private final Optional rackId; /** * The topicIds of topics that the member is subscribed to. */ - final Collection subscribedTopics; + private final Collection subscribedTopics; /** - * Partitions assigned for this member grouped by topicId + * Partitions assigned for this member keyed by topicId */ - final Map> currentAssignmentTopicIdPartitions; + private final Map> assignedTopicIdPartitions; - public AssignmentMemberSpec( + public Optional instanceId() { + return instanceId; + } + public Optional rackId() { + return rackId; + } + + public Collection subscribedTopics() { + return subscribedTopics; + } + + public Map> assignmentTopicIdPartitions() { + return assignedTopicIdPartitions; + } + + public AssignmentMemberSpec( Optional instanceId, Optional rackId, Collection subscribedTopics, - Map> currentAssignmentTopicIdPartitions + Map> assignedTopicIdPartitions ) { Objects.requireNonNull(instanceId); Objects.requireNonNull(rackId); Objects.requireNonNull(subscribedTopics); - Objects.requireNonNull(currentAssignmentTopicIdPartitions); + Objects.requireNonNull(assignedTopicIdPartitions); this.instanceId = instanceId; this.rackId = rackId; this.subscribedTopics = subscribedTopics; - this.currentAssignmentTopicIdPartitions = currentAssignmentTopicIdPartitions; + this.assignedTopicIdPartitions = assignedTopicIdPartitions; } @Override @@ -75,7 +90,7 @@ public boolean equals(Object o) { if (!instanceId.equals(that.instanceId)) return false; if (!rackId.equals(that.rackId)) return false; if (!subscribedTopics.equals(that.subscribedTopics)) return false; - return currentAssignmentTopicIdPartitions.equals(that.currentAssignmentTopicIdPartitions); + return assignedTopicIdPartitions.equals(that.assignedTopicIdPartitions); } @Override @@ -83,7 +98,7 @@ public int hashCode() { int result = instanceId.hashCode(); result = 31 * result + rackId.hashCode(); result = 31 * result + subscribedTopics.hashCode(); - result = 31 * result + currentAssignmentTopicIdPartitions.hashCode(); + result = 31 * result + assignedTopicIdPartitions.hashCode(); return result; } @@ -92,7 +107,7 @@ public String toString() { return "AssignmentMemberSpec(instanceId=" + instanceId + ", rackId=" + rackId + ", subscribedTopics=" + subscribedTopics + - ", targetPartitions=" + currentAssignmentTopicIdPartitions + + ", assignedTopicIdPartitions=" + assignedTopicIdPartitions + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java index f56b272b787c3..e36a66fd4ebd2 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java @@ -28,12 +28,12 @@ public class AssignmentSpec { /** * The members keyed by member id. */ - final Map members; + private final Map members; /** * The topics' metadata keyed by topic id */ - final Map topics; + private final Map topics; public AssignmentSpec( Map members, @@ -45,6 +45,14 @@ public AssignmentSpec( this.topics = topics; } + public Map members() { + return members; + } + + public Map topics() { + return topics; + } + @Override public boolean equals(Object o) { if (this == o) return true; diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java index ba79c82b5acd4..43f9a003efcf7 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java @@ -25,12 +25,12 @@ public class AssignmentTopicMetadata { /** * The topic name. */ - final String topicName; + private final String topicName; /** * The number of partitions. */ - final int numPartitions; + private final int numPartitions; public AssignmentTopicMetadata( String topicName, @@ -41,6 +41,14 @@ public AssignmentTopicMetadata( this.numPartitions = numPartitions; } + public String topicName() { + return topicName; + } + + public int numPartitions() { + return numPartitions; + } + @Override public boolean equals(Object o) { if (this == o) return true; diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java index 5c0199aaec5a5..70ff9bafcca3f 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java @@ -26,7 +26,7 @@ public class GroupAssignment { /** * The member assignments keyed by member id. */ - final Map members; + private final Map members; public GroupAssignment( Map members @@ -35,6 +35,10 @@ public GroupAssignment( this.members = members; } + public Map members() { + return members; + } + @Override public boolean equals(Object o) { if (this == o) return true; diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java index fa05e4398222b..09bf1789df982 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java @@ -27,17 +27,17 @@ */ public class MemberAssignment { /** - * The target partitions assigned to this member grouped by topicId. + * The target partitions assigned to this member keyed by topicId. */ - private final Map> topicIdPartitionsMap; + private final Map> assignedTopicIdPartitions; public MemberAssignment(Map> topicIdPartitionsForAssignment) { Objects.requireNonNull(topicIdPartitionsForAssignment); - this.topicIdPartitionsMap = topicIdPartitionsForAssignment; + this.assignedTopicIdPartitions = topicIdPartitionsForAssignment; } - public Map> getTopicIdPartitionsMap() { - return this.topicIdPartitionsMap; + public Map> assignedTopicIdPartitions() { + return this.assignedTopicIdPartitions; } @Override @@ -47,16 +47,16 @@ public boolean equals(Object o) { MemberAssignment that = (MemberAssignment) o; - return topicIdPartitionsMap.equals(that.topicIdPartitionsMap); + return assignedTopicIdPartitions.equals(that.assignedTopicIdPartitions); } @Override public int hashCode() { - return topicIdPartitionsMap.hashCode(); + return assignedTopicIdPartitions.hashCode(); } @Override public String toString() { - return "MemberAssignment ( Assignment per topic Id = " + topicIdPartitionsMap + ')'; + return "MemberAssignment (Assignment per topic Id = " + assignedTopicIdPartitions + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java index ea7d1114feb32..3aaccbcef0436 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java @@ -28,9 +28,9 @@ public class TopicIdToPartition { private final Optional> rackIds; public TopicIdToPartition(Uuid topicId, Integer topicPartition, Optional> rackIds) { - this.topicId = Objects.requireNonNull(topicId, "topicId can not be null"); - this.partition = Objects.requireNonNull(topicPartition, "topicPartition can not be null"); - this.rackIds = rackIds; + this.topicId = Objects.requireNonNull(topicId, "topicId cannot be null"); + this.partition = Objects.requireNonNull(topicPartition, "topicPartition cannot be null"); + this.rackIds = Objects.requireNonNull(rackIds, "rackId cannot be null"); } /** From 68d53b37b5e3fd78b8f96d7917b1fdc565e188dd Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Wed, 12 Apr 2023 13:00:07 -0700 Subject: [PATCH 11/44] Removed topicIdToPartition class, changed attribute names and added java doc for getter methods --- .../group/assignor/AssignmentMemberSpec.java | 52 ++++++++----- .../group/assignor/AssignmentSpec.java | 8 +- .../assignor/AssignmentTopicMetadata.java | 6 ++ .../group/assignor/GroupAssignment.java | 3 + .../group/assignor/MemberAssignment.java | 21 ++--- .../group/common/TopicIdToPartition.java | 76 ------------------- 6 files changed, 60 insertions(+), 106 deletions(-) delete mode 100644 group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java index 33dfe71165c85..8bcc08dc19eb3 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java @@ -39,45 +39,57 @@ public class AssignmentMemberSpec { private final Optional rackId; /** - * The topicIds of topics that the member is subscribed to. + * Topics Ids that the member is subscribed to. */ - private final Collection subscribedTopics; + private final Collection subscribedTopicIds; /** - * Partitions assigned for this member keyed by topicId + * Partitions assigned keyed by topicId. */ - private final Map> assignedTopicIdPartitions; + private final Map> assignedPartitions; + /** + * @return The instance ID as an Optional. + */ public Optional instanceId() { return instanceId; } + /** + * @return The rack ID as an Optional. + */ public Optional rackId() { return rackId; } - public Collection subscribedTopics() { - return subscribedTopics; + /** + * @return Collection of subscribed topic Ids. + */ + public Collection subscribedTopicIds() { + return subscribedTopicIds; } - public Map> assignmentTopicIdPartitions() { - return assignedTopicIdPartitions; + /** + * @return Assigned partitions keyed by topic Ids. + */ + public Map> assignedPartitions() { + return assignedPartitions; } public AssignmentMemberSpec( Optional instanceId, Optional rackId, - Collection subscribedTopics, - Map> assignedTopicIdPartitions + Collection subscribedTopicIds, + Map> assignedPartitions ) { Objects.requireNonNull(instanceId); Objects.requireNonNull(rackId); - Objects.requireNonNull(subscribedTopics); - Objects.requireNonNull(assignedTopicIdPartitions); + Objects.requireNonNull(subscribedTopicIds); + Objects.requireNonNull(assignedPartitions); this.instanceId = instanceId; this.rackId = rackId; - this.subscribedTopics = subscribedTopics; - this.assignedTopicIdPartitions = assignedTopicIdPartitions; + this.subscribedTopicIds = subscribedTopicIds; + this.assignedPartitions = assignedPartitions; } @Override @@ -89,16 +101,16 @@ public boolean equals(Object o) { if (!instanceId.equals(that.instanceId)) return false; if (!rackId.equals(that.rackId)) return false; - if (!subscribedTopics.equals(that.subscribedTopics)) return false; - return assignedTopicIdPartitions.equals(that.assignedTopicIdPartitions); + if (!subscribedTopicIds.equals(that.subscribedTopicIds)) return false; + return assignedPartitions.equals(that.assignedPartitions); } @Override public int hashCode() { int result = instanceId.hashCode(); result = 31 * result + rackId.hashCode(); - result = 31 * result + subscribedTopics.hashCode(); - result = 31 * result + assignedTopicIdPartitions.hashCode(); + result = 31 * result + subscribedTopicIds.hashCode(); + result = 31 * result + assignedPartitions.hashCode(); return result; } @@ -106,8 +118,8 @@ public int hashCode() { public String toString() { return "AssignmentMemberSpec(instanceId=" + instanceId + ", rackId=" + rackId + - ", subscribedTopics=" + subscribedTopics + - ", assignedTopicIdPartitions=" + assignedTopicIdPartitions + + ", subscribedTopicIds=" + subscribedTopicIds + + ", assignedPartitions=" + assignedPartitions + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java index e36a66fd4ebd2..39426dadfcce1 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java @@ -31,7 +31,7 @@ public class AssignmentSpec { private final Map members; /** - * The topics' metadata keyed by topic id + * The topics' metadata keyed by topic id. */ private final Map topics; @@ -45,10 +45,16 @@ public AssignmentSpec( this.topics = topics; } + /** + * @return Member metadata keyed by member Ids. + */ public Map members() { return members; } + /** + * @return Topic metadata keyed by topic Ids. + */ public Map topics() { return topics; } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java index 43f9a003efcf7..4fb9c63c02eaa 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java @@ -41,10 +41,16 @@ public AssignmentTopicMetadata( this.numPartitions = numPartitions; } + /** + * @return The topic name. + */ public String topicName() { return topicName; } + /** + * @return The number of partitions present for the topic. + */ public int numPartitions() { return numPartitions; } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java index 70ff9bafcca3f..a0464ab272eef 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java @@ -35,6 +35,9 @@ public GroupAssignment( this.members = members; } + /** + * @return Member assignments keyed by member Ids. + */ public Map members() { return members; } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java index 09bf1789df982..0a8dc204d2802 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java @@ -29,15 +29,18 @@ public class MemberAssignment { /** * The target partitions assigned to this member keyed by topicId. */ - private final Map> assignedTopicIdPartitions; + private final Map> targetPartitions; - public MemberAssignment(Map> topicIdPartitionsForAssignment) { - Objects.requireNonNull(topicIdPartitionsForAssignment); - this.assignedTopicIdPartitions = topicIdPartitionsForAssignment; + public MemberAssignment(Map> targetPartitions) { + Objects.requireNonNull(targetPartitions); + this.targetPartitions = targetPartitions; } - public Map> assignedTopicIdPartitions() { - return this.assignedTopicIdPartitions; + /** + * @return Target partition numbers keyed by topic Ids. + */ + public Map> targetPartitions() { + return this.targetPartitions; } @Override @@ -47,16 +50,16 @@ public boolean equals(Object o) { MemberAssignment that = (MemberAssignment) o; - return assignedTopicIdPartitions.equals(that.assignedTopicIdPartitions); + return targetPartitions.equals(that.targetPartitions); } @Override public int hashCode() { - return assignedTopicIdPartitions.hashCode(); + return targetPartitions.hashCode(); } @Override public String toString() { - return "MemberAssignment (Assignment per topic Id = " + assignedTopicIdPartitions + ')'; + return "MemberAssignment (Target partitions = " + targetPartitions + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java deleted file mode 100644 index 3aaccbcef0436..0000000000000 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/TopicIdToPartition.java +++ /dev/null @@ -1,76 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.kafka.coordinator.group.common; - -import org.apache.kafka.common.Uuid; - -import java.util.List; -import java.util.Objects; -import java.util.Optional; - -public class TopicIdToPartition { - private final Uuid topicId; - private final Integer partition; - private final Optional> rackIds; - - public TopicIdToPartition(Uuid topicId, Integer topicPartition, Optional> rackIds) { - this.topicId = Objects.requireNonNull(topicId, "topicId cannot be null"); - this.partition = Objects.requireNonNull(topicPartition, "topicPartition cannot be null"); - this.rackIds = Objects.requireNonNull(rackIds, "rackId cannot be null"); - } - - /** - * @return Universally unique id representing this topic partition. - */ - public Uuid topicId() { - return topicId; - } - - /** - * @return the partition number. - */ - public int partition() { - return partition; - } - - @Override - public boolean equals(Object o) { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - TopicIdToPartition that = (TopicIdToPartition) o; - return topicId.equals(that.topicId) && - partition.equals(that.partition); - } - - @Override - public int hashCode() { - final int prime = 31; - int result = prime + topicId.hashCode(); - result = prime * result + partition.hashCode(); - return result; - } - - @Override - public String toString() { - return topicId() + "-" + partition(); - } - -} From 48a982232c6f1ff371cb44333fd12d6b37ef381a Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 23 Mar 2023 13:02:31 -0700 Subject: [PATCH 12/44] Server Side Sticky Range Assignor full implementation --- checkstyle/suppressions.xml | 8 + .../ServerSideStickyRangeAssignor.java | 268 +++++++++++ .../ServerSideStickyRangeAssignorTest.java | 428 ++++++++++++++++++ 3 files changed, 704 insertions(+) create mode 100644 group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index 6962d0a5c4abe..1bf330e52fc73 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -318,6 +318,14 @@ + + + + + diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java new file mode 100644 index 0000000000000..186f8af1f0220 --- /dev/null +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java @@ -0,0 +1,268 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + *

The Server Side Sticky Range Assignor inherits properties of both the range assignor and the sticky assignor. + * Properties are :- + *

    + *
  • 1) Each consumer must get at least one partition per topic that it is subscribed to whenever the number of consumers is + * less than or equal to the number of partitions for that topic. (Range)
    • + *
  • 2) Partitions should be assigned to consumers in a way that facilitates join operations where required. (Range)
    • + * This can only be done if the topics are co-partitioned in the first place + * Co-partitioned:- + * Two streams are co-partitioned if the following conditions are met:- + * ->The keys must have the same schemas + * ->The topics involved must have the same number of partitions + *
  • 3) Consumers should retain as much as their previous assignment as possible. (Sticky)
    • + *
+ *

+ * + *

The algorithm works mainly in 5 steps described below + *

    + *
  • 1) Get a map of the consumersPerTopic created using the member subscriptions.
    • + *
  • 2) Get a list of consumers (potentiallyUnfilled) that have not met the minimum required quota for assignment AND + * get a list of sticky partitions that we want to retain in the new assignment.
    • + *
  • 3) Add consumers from potentiallyUnfilled to Unfilled if they haven't met the total required quota = minQuota + (if necessary) extraPartition
    • + *
  • 4) Get a list of available partitions by calculating the difference between total partitions and assigned sticky partitions
    • + *
  • 5) Iterate through unfilled consumers and assign partitions from available partitions
    • + *
+ *

+ * + */ +package org.apache.kafka.coordinator.group.assignor; + + +import org.apache.kafka.common.Uuid; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.HashMap; +import java.util.Collections; + +import static java.lang.Math.min; + +public class ServerSideStickyRangeAssignor implements PartitionAssignor { + + public static final String RANGE_ASSIGNOR_NAME = "range-sticky"; + + @Override + public String name() { + return RANGE_ASSIGNOR_NAME; + } + + protected static void putList(Map> map, K key, V value) { + List list = map.computeIfAbsent(key, k -> new ArrayList<>()); + list.add(value); + } + + protected static void putSet(Map> map, K key, V value) { + Set set = map.computeIfAbsent(key, k -> new HashSet<>()); + set.add(value); + } + + static class Pair { + private final T first; + private final U second; + + public Pair(T first, U second) { + this.first = first; + this.second = second; + } + + public T getFirst() { + return first; + } + + public U getSecond() { + return second; + } + + @Override + public String toString() { + return "(" + first + ", " + second + ")"; + } + } + + // Returns a map of the list of consumers per Topic (keyed by topicId) + private Map> consumersPerTopic(AssignmentSpec assignmentSpec) { + Map> mapTopicsToConsumers = new HashMap<>(); + Map membersData = assignmentSpec.members; + + for (Map.Entry memberEntry : membersData.entrySet()) { + String memberId = memberEntry.getKey(); + AssignmentMemberSpec memberMetadata = memberEntry.getValue(); + List topics = new ArrayList<>(memberMetadata.subscribedTopics); + for (Uuid topicId: topics) { + putList(mapTopicsToConsumers, topicId, memberId); + } + } + return mapTopicsToConsumers; + } + + private Map> getAvailablePartitionsPerTopic(AssignmentSpec assignmentSpec, Map> assignedStickyPartitionsPerTopic) { + Map> availablePartitionsPerTopic = new HashMap<>(); + Map topicsMetadata = assignmentSpec.topics; + // Iterate through the topics map provided in assignmentSpec + for (Map.Entry topicMetadataEntry : topicsMetadata.entrySet()) { + Uuid topicId = topicMetadataEntry.getKey(); + availablePartitionsPerTopic.put(topicId, new ArrayList<>()); + int numPartitions = topicsMetadata.get(topicId).numPartitions; + // since the loop iterates from 0 to n, the partitions will be in ascending order within the list of available partitions per topic + Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.get(topicId); + for (int i = 0; i < numPartitions; i++) { + if (assignedStickyPartitionsForTopic == null || !assignedStickyPartitionsForTopic.contains(i)) { + availablePartitionsPerTopic.get(topicId).add(i); + } + } + } + return availablePartitionsPerTopic; + } + + @Override + public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAssignorException { + Map>> membersWithNewAssignmentPerTopic = new HashMap<>(); + // Step 1 + Map> consumersPerTopic = consumersPerTopic(assignmentSpec); + // Step 2 + Map>> unfilledConsumersPerTopic = new HashMap<>(); + Map> assignedStickyPartitionsPerTopic = new HashMap<>(); + + for (Map.Entry> topicEntry : consumersPerTopic.entrySet()) { + Uuid topicId = topicEntry.getKey(); + // For each topic we have a temporary list of consumers stored in potentiallyUnfilledConsumers. + // The list is populated with consumers that satisfy one of the two conditions :- + // 1) Consumers that have the minimum required number of partitions .i.e numPartitionsPerConsumer BUT they could be assigned an extra partition later on. + // In this case we add the consumer to the unfilled consumers map iff an extra partition needs to be assigned to it. + // 2) Consumers that don't have the minimum required partitions, so irrespective of whether they get an extra partition or not they get added to the unfilled map later. + List> potentiallyUnfilledConsumers = new ArrayList<>(); + List consumersForTopic = topicEntry.getValue(); + + AssignmentTopicMetadata topicData = assignmentSpec.topics.get(topicId); + int numPartitionsForTopic = topicData.numPartitions; + int numPartitionsPerConsumer = numPartitionsForTopic / consumersForTopic.size(); + // Each consumer can get only one extra partition per topic after receiving the minimum quota = numPartitionsPerConsumer + int numConsumersWithExtraPartition = numPartitionsForTopic % consumersForTopic.size(); + + for (String memberId: consumersForTopic) { + // Size of the older assignment, this will be 0 when assign is called for the first time. + // The older assignment is required when a reassignment occurs to ensure stickiness. + int currentAssignmentSize = 0; + List currentAssignmentListForTopic = new ArrayList<>(); + // Convert the set to a list first and sort the partitions in numeric order since we want the same partition numbers from each topic + // to go to the same consumer in case of co-partitioned topics. + Set currentAssignmentSetForTopic = assignmentSpec.members.get(memberId).currentAssignmentPerTopic.get(topicId); + // We want to make sure that the currentAssignment is not null to avoid a null pointer exception + if (currentAssignmentSetForTopic != null) { + currentAssignmentListForTopic = new ArrayList<>(currentAssignmentSetForTopic); + currentAssignmentSize = currentAssignmentListForTopic.size(); + } + // Initially, minRequiredQuota is the minimum number of partitions that each consumer should get. + // The value is at least 1 unless numConsumers subscribed is greater than numPartitions. In such cases, all consumers get assigned "extra partitions". + int minRequiredQuota = numPartitionsPerConsumer; + // We need to make sure that the partitions we're keeping are still part of the topic metadata. + // Ex:- In case the number of partitions is reduced for a topic, the sticky partitions must be part of the new set of partitions available. + for (int i = 0; i < currentAssignmentListForTopic.size() && currentAssignmentSize > 0; i++) { + // numPartitionsForTopic - 1 is the max possible partition number + if (currentAssignmentListForTopic.get(i) >= numPartitionsForTopic) { + currentAssignmentSize--; + } + } + // If there are previously assigned partitions present, we want to retain + if (currentAssignmentSize > 0) { + // We either need to retain currentSize number of partitions when currentSize < required OR required number of partitions otherwise. + int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); + for (int i = 0; i < retainedPartitionsCount; i++) { + Collections.sort(currentAssignmentListForTopic); + putSet(assignedStickyPartitionsPerTopic, topicId, currentAssignmentListForTopic.get(i)); + membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); + } + } + // Number of partitions left to reach the minRequiredQuota + int remaining = minRequiredQuota - currentAssignmentSize; + + // There are 3 cases w.r.t value of remaining + // 1) remaining < 0 this means that the consumer has more than the min required amount. + // It could have an extra partition, so we check for that. + if (remaining < 0 && numConsumersWithExtraPartition > 0) { + // In order to remain as sticky as possible, since the order of members can be different, we want the consumers that already had extra + // partitions to retain them if it's still required, instead of assigning the extras to the first few consumers directly. + // Ex:- If two consumers out of 3 are supposed to get an extra partition and currently 1 of them already has the extra, we want this consumer + // to retain it first and later if we have partitions left they will be assigned to the first few consumers and the unfilled map is updated + numConsumersWithExtraPartition--; + // Since we already added the minimumRequiredQuota of partitions in the previous step (until minReq - 1), we just need to + // add the extra partition which will be present at the index right after min quota is satisfied. + putSet(assignedStickyPartitionsPerTopic, topicId, currentAssignmentListForTopic.get(minRequiredQuota)); + membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(minRequiredQuota)); + } else { + // 3) If remaining = 0 it has min req partitions but there is scope for getting an extra partition later on, so it is a potentialUnfilledConsumer. + // 4) If remaining > 0 it doesn't even have the min required partitions, and it definitely is unfilled so it should be added to potentialUnfilledConsumers. + Pair newPair = new Pair<>(memberId, remaining); + potentiallyUnfilledConsumers.add(newPair); + } + } + + // Step 3 + // Iterate through potentially unfilled consumers and if remaining > 0 after considering the extra partitions assignment, add to the unfilled list. + for (Pair pair : potentiallyUnfilledConsumers) { + String memberId = pair.getFirst(); + Integer remaining = pair.getSecond(); + if (numConsumersWithExtraPartition > 0) { + remaining++; + numConsumersWithExtraPartition--; + } + if (remaining > 0) { + Pair newPair = new Pair<>(memberId, remaining); + putList(unfilledConsumersPerTopic, topicId, newPair); + } + } + } + + // Step 4 + // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get available partitions. + Map> availablePartitionsPerTopic = getAvailablePartitionsPerTopic(assignmentSpec, assignedStickyPartitionsPerTopic); + + // Step 5 + // Iterate through the unfilled consumers list and assign remaining number of partitions from the availablePartitions list + for (Map.Entry>> unfilledEntry : unfilledConsumersPerTopic.entrySet()) { + Uuid topicId = unfilledEntry.getKey(); + List> unfilledConsumersForTopic = unfilledEntry.getValue(); + + for (Pair stringIntegerPair : unfilledConsumersForTopic) { + String memberId = stringIntegerPair.getFirst(); + int remaining = stringIntegerPair.getSecond(); + // assign the first few partitions from the list and then delete them from the availablePartitions list for this topic + List subList = availablePartitionsPerTopic.get(topicId).subList(0, remaining); + membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).addAll(subList); + availablePartitionsPerTopic.get(topicId).removeAll(subList); + } + } + + // Consolidate the maps into MemberAssignment and then finally map each consumer to a MemberAssignment. + Map membersWithNewAssignment = new HashMap<>(); + for (Map.Entry>> consumer : membersWithNewAssignmentPerTopic.entrySet()) { + String consumerId = consumer.getKey(); + Map> assignmentPerTopic = consumer.getValue(); + membersWithNewAssignment.computeIfAbsent(consumerId, k -> new MemberAssignment(assignmentPerTopic)); + } + + return new GroupAssignment(membersWithNewAssignment); + } +} + diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java new file mode 100644 index 0000000000000..228157280f610 --- /dev/null +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java @@ -0,0 +1,428 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.coordinator.group.assignor; + + +import org.apache.kafka.common.Uuid; + +import org.junit.jupiter.api.Test; + +import java.util.Map; +import java.util.Set; + + + +import java.util.*; + +import static org.junit.jupiter.api.Assertions.*; + +public class ServerSideStickyRangeAssignorTest { + + private final ServerSideStickyRangeAssignor assignor = new ServerSideStickyRangeAssignor(); + + private final String topic1Name = "topic1"; + private final Uuid topic1Uuid = Uuid.randomUuid(); + + private final String topic2Name = "topic2"; + private final Uuid topic2Uuid = Uuid.randomUuid(); + + private final String topic3Name = "topic3"; + private final Uuid topic3Uuid = Uuid.randomUuid(); + + private final String consumerA = "A"; + private final String consumerB = "B"; + private final String consumerC = "C"; + + @Test + public void testOneConsumerNoTopic() { + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + Map members = new HashMap<>(); + List subscribedTopics = new ArrayList<>(); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopics, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment groupAssignment = assignor.assign(assignmentSpec); + + assertTrue(groupAssignment.getMembers().isEmpty()); + } + + @Test + public void testFirstAssignmentTwoConsumersTwoTopicsSameSubscriptions() { + // A -> T1, T3 // B -> T1, T3 // T1 -> 3 Partitions // T3 -> 2 Partitions + // Topics + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + // Members + Map members = new HashMap<>(); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + // Topic 3 Partitions Assignment + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + assertAssignment(expectedAssignment, computedAssignment); + } + + @Test + public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() { + // A -> T1, T2 // B -> T3 // C -> T2, T3 // T1 -> 3 Partitions // T2 -> 3 Partitions // T3 -> 2 Partitions + // Topics + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + // Members + Map members = new HashMap<>(); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); + // Consumer B + List subscribedTopicsB = new ArrayList<>(Collections.singletonList(topic3Uuid)); + members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); + // Consumer B + List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic2Uuid, topic3Uuid)); + members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + // Topic 3 Partitions Assignment + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + assertAssignment(expectedAssignment, computedAssignment); + } + + @Test + public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { + // Topics + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + // Topic 3 has 2 partitions but three consumers subscribed to it - one of them will not get an assignment + topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + // Members + Map members = new HashMap<>(); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); + // Consumer C + List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + assertAssignment(expectedAssignment, computedAssignment); + } + + @Test + public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerAdded() { + // Topics + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 2)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 2)); + // Members + Map members = new HashMap<>(); + // Add a new consumer to trigger a re-assignment + // Consumer C + List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Collections.singletonList(0))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Collections.singletonList(0))); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(1))); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(1))); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + // Consumer C shouldn't get any assignment, due to stickiness A, B retain their assignments + assertNull(computedAssignment.getMembers().get(consumerC)); + assertAssignment(expectedAssignment, computedAssignment); + assertCoPartitionJoinProperty(computedAssignment); + } + + @Test + public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { + // T1, T2 both have 3 partitions each when first assignment was calculated -> currentAssignmentForX + Map members = new HashMap<>(); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + + // Simulating adding a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 4)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 4)); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(2, 3))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(2, 3))); + + assertAssignment(expectedAssignment, computedAssignment); + assertCoPartitionJoinProperty(computedAssignment); + } + + @Test + public void testReassignmentWhenOnePartitionRemovedForTwoConsumersTwoTopics() { + Map members = new HashMap<>(); + // T1, T2 both have 3 partitions each initially and after removal have 2 partitions each + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + + // Remove a partition from both topic1 and topic 2 -> total partitions = 2 + // Simulating removing a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 2)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 2)); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + assertAssignment(expectedAssignment, computedAssignment); + assertCoPartitionJoinProperty(computedAssignment); + } + + @Test + public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoConsumersTwoTopics() { + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); + + Map members = new HashMap<>(); + // Consumer A + List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + // Add a new consumer to trigger a re-assignment + // Consumer C + List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + + assertAssignment(expectedAssignment, computedAssignment); + assertCoPartitionJoinProperty(computedAssignment); + } + + @Test + public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoConsumersTwoTopics() { + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); + + Map members = new HashMap<>(); + // Consumer A was removed + + // Consumer B + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); + + assertAssignment(expectedAssignment, computedAssignment); + assertCoPartitionJoinProperty(computedAssignment); + } + + @Test + public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignmentWithThreeConsumersTwoTopics() { + + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); + topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + + // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3 + // Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2 + Map members = new HashMap<>(); + // Consumer A + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1, 2))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Collections.singletonList(0))); + // Change subscriptions + List subscribedTopicsA = new ArrayList<>(Collections.singletonList(topic1Uuid)); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + // Consumer B + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(1))); + // Change subscriptions + List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid, topic3Uuid)); + members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + // Consumer C + Map> currentAssignmentForC = new HashMap<>(); + currentAssignmentForC.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForC.put(topic3Uuid, new HashSet<>(Arrays.asList(0, 1))); + // Change subscriptions + List subscribedTopicsC = new ArrayList<>(Collections.singletonList(topic2Uuid)); + members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + // Topic 2 Partitions Assignment + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + // Topic 3 Partitions Assignment + expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + + assertAssignment(expectedAssignment, computedAssignment); + } + + // We have a set of sets with the partitions that should be distributed amongst the consumers, if it exists then remove it from the set. + private void assertAssignment(Map>> expectedAssignment, GroupAssignment computedGroupAssignment) { + for (Map.Entry member : computedGroupAssignment.getMembers().entrySet()) { + Map> computedAssignmentForMember = member.getValue().getAssignmentPerTopic(); + for (Map.Entry> assignmentForTopic : computedAssignmentForMember.entrySet()) { + Uuid topicId = assignmentForTopic.getKey(); + Set assignmentPartitionsSet = assignmentForTopic.getValue(); + assertTrue(expectedAssignment.get(topicId).contains(assignmentPartitionsSet)); + expectedAssignment.remove(assignmentPartitionsSet); + } + } + } + + private void assertCoPartitionJoinProperty(GroupAssignment groupAssignment) { + for (Map.Entry member : groupAssignment.getMembers().entrySet()) { + Map> computedAssignmentForMember = member.getValue().getAssignmentPerTopic(); + Set compareSet = new HashSet<>(); + for (Map.Entry> topicAssignment : computedAssignmentForMember.entrySet()) { + Set partitionsForTopicSet = topicAssignment.getValue(); + if (compareSet.isEmpty()) { + compareSet = partitionsForTopicSet; + } + assertEquals(compareSet, partitionsForTopicSet); + } + } + } +} From dcb8198355e82f36f12d042ed0d0a4c1d71ea8c6 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Mon, 27 Mar 2023 20:17:14 -0700 Subject: [PATCH 13/44] removed reduce partitions case --- .../ServerSideStickyRangeAssignor.java | 19 +++------ .../ServerSideStickyRangeAssignorTest.java | 39 ------------------- 2 files changed, 6 insertions(+), 52 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java index 186f8af1f0220..e9d90ccfb7848 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java @@ -108,7 +108,7 @@ private Map> consumersPerTopic(AssignmentSpec assignmentSpec) for (Map.Entry memberEntry : membersData.entrySet()) { String memberId = memberEntry.getKey(); AssignmentMemberSpec memberMetadata = memberEntry.getValue(); - List topics = new ArrayList<>(memberMetadata.subscribedTopics); + List topics = memberMetadata.subscribedTopics; for (Uuid topicId: topics) { putList(mapTopicsToConsumers, topicId, memberId); } @@ -176,20 +176,13 @@ public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAss // Initially, minRequiredQuota is the minimum number of partitions that each consumer should get. // The value is at least 1 unless numConsumers subscribed is greater than numPartitions. In such cases, all consumers get assigned "extra partitions". int minRequiredQuota = numPartitionsPerConsumer; - // We need to make sure that the partitions we're keeping are still part of the topic metadata. - // Ex:- In case the number of partitions is reduced for a topic, the sticky partitions must be part of the new set of partitions available. - for (int i = 0; i < currentAssignmentListForTopic.size() && currentAssignmentSize > 0; i++) { - // numPartitionsForTopic - 1 is the max possible partition number - if (currentAssignmentListForTopic.get(i) >= numPartitionsForTopic) { - currentAssignmentSize--; - } - } + // If there are previously assigned partitions present, we want to retain if (currentAssignmentSize > 0) { // We either need to retain currentSize number of partitions when currentSize < required OR required number of partitions otherwise. int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); + Collections.sort(currentAssignmentListForTopic); for (int i = 0; i < retainedPartitionsCount; i++) { - Collections.sort(currentAssignmentListForTopic); putSet(assignedStickyPartitionsPerTopic, topicId, currentAssignmentListForTopic.get(i)); membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); } @@ -243,14 +236,14 @@ public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAss for (Map.Entry>> unfilledEntry : unfilledConsumersPerTopic.entrySet()) { Uuid topicId = unfilledEntry.getKey(); List> unfilledConsumersForTopic = unfilledEntry.getValue(); - + int newStartPointer = 0; for (Pair stringIntegerPair : unfilledConsumersForTopic) { String memberId = stringIntegerPair.getFirst(); int remaining = stringIntegerPair.getSecond(); // assign the first few partitions from the list and then delete them from the availablePartitions list for this topic - List subList = availablePartitionsPerTopic.get(topicId).subList(0, remaining); + List subList = availablePartitionsPerTopic.get(topicId).subList(newStartPointer, newStartPointer + remaining); + newStartPointer += remaining; membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).addAll(subList); - availablePartitionsPerTopic.get(topicId).removeAll(subList); } } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java index 228157280f610..54a73229718bb 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java @@ -242,45 +242,6 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { assertCoPartitionJoinProperty(computedAssignment); } - @Test - public void testReassignmentWhenOnePartitionRemovedForTwoConsumersTwoTopics() { - Map members = new HashMap<>(); - // T1, T2 both have 3 partitions each initially and after removal have 2 partitions each - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); - - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); - - // Remove a partition from both topic1 and topic 2 -> total partitions = 2 - // Simulating removing a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 2)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 2)); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - - Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - - assertAssignment(expectedAssignment, computedAssignment); - assertCoPartitionJoinProperty(computedAssignment); - } - @Test public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoConsumersTwoTopics() { Map topics = new HashMap<>(); From 9f3a423d6a2b482c6ac6d74c4e65bbed70898dfa Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 28 Mar 2023 12:21:33 -0700 Subject: [PATCH 14/44] removed reduce partitions case --- .../group/assignor/ServerSideStickyRangeAssignor.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java index e9d90ccfb7848..8c5f021405bde 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java @@ -60,7 +60,7 @@ public class ServerSideStickyRangeAssignor implements PartitionAssignor { - public static final String RANGE_ASSIGNOR_NAME = "range-sticky"; + public static final String RANGE_ASSIGNOR_NAME = "range"; @Override public String name() { From a925b02dd3d78c8b9c1415c083f0544edcb404fe Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Wed, 29 Mar 2023 12:30:13 -0700 Subject: [PATCH 15/44] Addressed PR comments --- .../ServerSideStickyRangeAssignor.java | 72 +++++++++---------- 1 file changed, 34 insertions(+), 38 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java index 8c5f021405bde..c2e2dcae54c5c 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java @@ -14,6 +14,20 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +package org.apache.kafka.coordinator.group.assignor; + + +import org.apache.kafka.common.Uuid; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.HashMap; +import java.util.Collections; + +import static java.lang.Math.min; /** *

The Server Side Sticky Range Assignor inherits properties of both the range assignor and the sticky assignor. @@ -43,21 +57,6 @@ *

* */ -package org.apache.kafka.coordinator.group.assignor; - - -import org.apache.kafka.common.Uuid; - -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.HashMap; -import java.util.Collections; - -import static java.lang.Math.min; - public class ServerSideStickyRangeAssignor implements PartitionAssignor { public static final String RANGE_ASSIGNOR_NAME = "range"; @@ -67,12 +66,12 @@ public String name() { return RANGE_ASSIGNOR_NAME; } - protected static void putList(Map> map, K key, V value) { + private static void putList(Map> map, K key, V value) { List list = map.computeIfAbsent(key, k -> new ArrayList<>()); list.add(value); } - protected static void putSet(Map> map, K key, V value) { + private static void putSet(Map> map, K key, V value) { Set set = map.computeIfAbsent(key, k -> new HashSet<>()); set.add(value); } @@ -119,18 +118,19 @@ private Map> consumersPerTopic(AssignmentSpec assignmentSpec) private Map> getAvailablePartitionsPerTopic(AssignmentSpec assignmentSpec, Map> assignedStickyPartitionsPerTopic) { Map> availablePartitionsPerTopic = new HashMap<>(); Map topicsMetadata = assignmentSpec.topics; - // Iterate through the topics map provided in assignmentSpec + for (Map.Entry topicMetadataEntry : topicsMetadata.entrySet()) { Uuid topicId = topicMetadataEntry.getKey(); - availablePartitionsPerTopic.put(topicId, new ArrayList<>()); + ArrayList availablePartitionsForTopic = new ArrayList<>(); int numPartitions = topicsMetadata.get(topicId).numPartitions; // since the loop iterates from 0 to n, the partitions will be in ascending order within the list of available partitions per topic - Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.get(topicId); + Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.getOrDefault(topicId, new HashSet<>()); for (int i = 0; i < numPartitions; i++) { - if (assignedStickyPartitionsForTopic == null || !assignedStickyPartitionsForTopic.contains(i)) { - availablePartitionsPerTopic.get(topicId).add(i); + if (!assignedStickyPartitionsForTopic.contains(i)) { + availablePartitionsForTopic.add(i); } } + availablePartitionsPerTopic.put(topicId, availablePartitionsForTopic); } return availablePartitionsPerTopic; } @@ -156,28 +156,23 @@ public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAss AssignmentTopicMetadata topicData = assignmentSpec.topics.get(topicId); int numPartitionsForTopic = topicData.numPartitions; - int numPartitionsPerConsumer = numPartitionsForTopic / consumersForTopic.size(); + // Initially, minRequiredQuota is the minimum number of partitions that each consumer should get i.e numPartitionsPerConsumer. + // Idle consumers case :- The numConsumers subscribed to a topic is greater than numPartitions. In such cases, all consumers get assigned via the "extra partitions" logic since min = 0. + int minRequiredQuota = numPartitionsForTopic / consumersForTopic.size(); // Each consumer can get only one extra partition per topic after receiving the minimum quota = numPartitionsPerConsumer int numConsumersWithExtraPartition = numPartitionsForTopic % consumersForTopic.size(); for (String memberId: consumersForTopic) { - // Size of the older assignment, this will be 0 when assign is called for the first time. - // The older assignment is required when a reassignment occurs to ensure stickiness. - int currentAssignmentSize = 0; - List currentAssignmentListForTopic = new ArrayList<>(); + // Convert the set to a list first and sort the partitions in numeric order since we want the same partition numbers from each topic // to go to the same consumer in case of co-partitioned topics. - Set currentAssignmentSetForTopic = assignmentSpec.members.get(memberId).currentAssignmentPerTopic.get(topicId); - // We want to make sure that the currentAssignment is not null to avoid a null pointer exception - if (currentAssignmentSetForTopic != null) { - currentAssignmentListForTopic = new ArrayList<>(currentAssignmentSetForTopic); - currentAssignmentSize = currentAssignmentListForTopic.size(); - } - // Initially, minRequiredQuota is the minimum number of partitions that each consumer should get. - // The value is at least 1 unless numConsumers subscribed is greater than numPartitions. In such cases, all consumers get assigned "extra partitions". - int minRequiredQuota = numPartitionsPerConsumer; + Set currentAssignmentSetForTopic = assignmentSpec.members.get(memberId).currentAssignmentPerTopic.getOrDefault(topicId, new HashSet<>()); + // Size of the older assignment, this will be 0 when assign is called for the first time. + // The older assignment is required when a reassignment occurs to ensure stickiness. + int currentAssignmentSize = currentAssignmentSetForTopic.size(); + List currentAssignmentListForTopic = new ArrayList<>(currentAssignmentSetForTopic); - // If there are previously assigned partitions present, we want to retain + // If there are previously assigned partitions present, we want to retain them. if (currentAssignmentSize > 0) { // We either need to retain currentSize number of partitions when currentSize < required OR required number of partitions otherwise. int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); @@ -187,7 +182,8 @@ public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAss membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); } } - // Number of partitions left to reach the minRequiredQuota + + // Number of partitions left to reach the minRequiredQuota. int remaining = minRequiredQuota - currentAssignmentSize; // There are 3 cases w.r.t value of remaining From 7d1626990abc6bfdf29b5c227751d10d8211ed5f Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Mon, 3 Apr 2023 11:54:59 -0700 Subject: [PATCH 16/44] Addressed PR comments --- .../assignor/ServerSideStickyRangeAssignorTest.java | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java index 54a73229718bb..5704d68cc5fde 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java @@ -108,7 +108,7 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() // Consumer B List subscribedTopicsB = new ArrayList<>(Collections.singletonList(topic3Uuid)); members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); - // Consumer B + // Consumer C List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic2Uuid, topic3Uuid)); members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); @@ -362,8 +362,8 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme // We have a set of sets with the partitions that should be distributed amongst the consumers, if it exists then remove it from the set. private void assertAssignment(Map>> expectedAssignment, GroupAssignment computedGroupAssignment) { - for (Map.Entry member : computedGroupAssignment.getMembers().entrySet()) { - Map> computedAssignmentForMember = member.getValue().getAssignmentPerTopic(); + for (MemberAssignment member : computedGroupAssignment.getMembers().values()) { + Map> computedAssignmentForMember = member.getAssignmentPerTopic(); for (Map.Entry> assignmentForTopic : computedAssignmentForMember.entrySet()) { Uuid topicId = assignmentForTopic.getKey(); Set assignmentPartitionsSet = assignmentForTopic.getValue(); @@ -374,11 +374,10 @@ private void assertAssignment(Map>> expectedAssignment, G } private void assertCoPartitionJoinProperty(GroupAssignment groupAssignment) { - for (Map.Entry member : groupAssignment.getMembers().entrySet()) { - Map> computedAssignmentForMember = member.getValue().getAssignmentPerTopic(); + for (MemberAssignment member : groupAssignment.getMembers().values()) { + Map> computedAssignmentForMember = member.getAssignmentPerTopic(); Set compareSet = new HashSet<>(); - for (Map.Entry> topicAssignment : computedAssignmentForMember.entrySet()) { - Set partitionsForTopicSet = topicAssignment.getValue(); + for (Set partitionsForTopicSet : computedAssignmentForMember.values()) { if (compareSet.isEmpty()) { compareSet = partitionsForTopicSet; } From 7369db3744b4c89f625f6b73db331eaf0512c4eb Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 6 Apr 2023 10:55:06 -0700 Subject: [PATCH 17/44] minor --- .../group/assignor/ServerSideStickyRangeAssignorTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java index 5704d68cc5fde..73a7860239e0a 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java @@ -98,7 +98,7 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() // Topics Map topics = new HashMap<>(); topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); // Members Map members = new HashMap<>(); From 918d262ae52892c13766649db847ede6928921f8 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 6 Apr 2023 11:12:16 -0700 Subject: [PATCH 18/44] Made subscribed topics a collection, removed * import exception --- checkstyle/suppressions.xml | 2 -- .../assignor/ServerSideStickyRangeAssignor.java | 12 ++++++------ .../ServerSideStickyRangeAssignorTest.java | 16 ++++++++++------ 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index 1bf330e52fc73..e2a501ae3a575 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -323,8 +323,6 @@ files="ServerSideStickyRangeAssignor.java"/> - > consumersPerTopic(AssignmentSpec assignmentSpec) for (Map.Entry memberEntry : membersData.entrySet()) { String memberId = memberEntry.getKey(); AssignmentMemberSpec memberMetadata = memberEntry.getValue(); - List topics = memberMetadata.subscribedTopics; + Collection topics = memberMetadata.subscribedTopics; for (Uuid topicId: topics) { putList(mapTopicsToConsumers, topicId, memberId); } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java index 73a7860239e0a..52d36eece5610 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java @@ -17,19 +17,23 @@ package org.apache.kafka.coordinator.group.assignor; - import org.apache.kafka.common.Uuid; import org.junit.jupiter.api.Test; +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.Set; - - -import java.util.*; - -import static org.junit.jupiter.api.Assertions.*; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; public class ServerSideStickyRangeAssignorTest { From 288fa1f5f853a0aa5d140bc930d2a8a52eeeff71 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Wed, 12 Apr 2023 11:28:40 -0700 Subject: [PATCH 19/44] Addressed some PR comments --- ...ckyRangeAssignor.java => RangeAssignor.java} | 17 ++++++++--------- ...AssignorTest.java => RangeAssignorTest.java} | 4 ++-- 2 files changed, 10 insertions(+), 11 deletions(-) rename group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/{ServerSideStickyRangeAssignor.java => RangeAssignor.java} (94%) rename group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/{ServerSideStickyRangeAssignorTest.java => RangeAssignorTest.java} (99%) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java similarity index 94% rename from group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java rename to group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index bb89296d0c0ee..6f613f17b65c0 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -36,9 +36,8 @@ *
  • 1) Each consumer must get at least one partition per topic that it is subscribed to whenever the number of consumers is * less than or equal to the number of partitions for that topic. (Range)
    • *
  • 2) Partitions should be assigned to consumers in a way that facilitates join operations where required. (Range)
    • - * This can only be done if the topics are co-partitioned in the first place - * Co-partitioned:- - * Two streams are co-partitioned if the following conditions are met:- + * This can only be done if every consumer is subscribed to the same topics and the topics are co-partitioned. + * Two streams are co-partitioned if the following conditions are met: * ->The keys must have the same schemas * ->The topics involved must have the same number of partitions *
  • 3) Consumers should retain as much as their previous assignment as possible. (Sticky)
    • @@ -47,17 +46,17 @@ * *

    The algorithm works mainly in 5 steps described below *

      - *
    • 1) Get a map of the consumersPerTopic created using the member subscriptions.
      • - *
    • 2) Get a list of consumers (potentiallyUnfilled) that have not met the minimum required quota for assignment AND + *
    • 1) Generate a map of consumersPerTopic using the member subscriptions.
      • + *
    • 2) Generate a list of consumers (potentiallyUnfilledConsumers) that have not met the minimum required quota for assignment AND * get a list of sticky partitions that we want to retain in the new assignment.
      • - *
    • 3) Add consumers from potentiallyUnfilled to Unfilled if they haven't met the total required quota = minQuota + (if necessary) extraPartition
      • - *
    • 4) Get a list of available partitions by calculating the difference between total partitions and assigned sticky partitions
      • - *
    • 5) Iterate through unfilled consumers and assign partitions from available partitions
      • + *
    • 3) Add consumers from the potentiallyUnfilled list to the Unfilled list if they haven't met the total required quota i.e. minimum number of partitions per member + 1 (if member is designated to receive one of the excess partitions)
      • + *
    • 4) Generate a list of unassigned partitions by calculating the difference between total partitions and already assigned (sticky) partitions
      • + *
    • 5) Iterate through unfilled consumers and assign partitions from the available partitions
      • *
    *

    * */ -public class ServerSideStickyRangeAssignor implements PartitionAssignor { +public class RangeAssignor implements PartitionAssignor { public static final String RANGE_ASSIGNOR_NAME = "range"; diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java similarity index 99% rename from group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java rename to group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 52d36eece5610..38224ff03df26 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -35,9 +35,9 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; -public class ServerSideStickyRangeAssignorTest { +public class RangeAssignorTest { - private final ServerSideStickyRangeAssignor assignor = new ServerSideStickyRangeAssignor(); + private final RangeAssignor assignor = new RangeAssignor(); private final String topic1Name = "topic1"; private final Uuid topic1Uuid = Uuid.randomUuid(); From a1bdb58e082d98b73f50860db5b35684f4eba51b Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Wed, 12 Apr 2023 20:01:59 -0700 Subject: [PATCH 20/44] Removed putList, putSet, changed generic pair to remainingAssignmentsForMember, addressed PR comments, changed Map.Entry to Map.forEach etc. --- .../group/assignor/RangeAssignor.java | 259 ++++++++---------- .../group/assignor/RangeAssignorTest.java | 12 +- 2 files changed, 123 insertions(+), 148 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 6f613f17b65c0..b44556ae1fce6 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -31,28 +31,30 @@ /** *

    The Server Side Sticky Range Assignor inherits properties of both the range assignor and the sticky assignor. - * Properties are :- - *

      - *
    • 1) Each consumer must get at least one partition per topic that it is subscribed to whenever the number of consumers is - * less than or equal to the number of partitions for that topic. (Range)
      • - *
    • 2) Partitions should be assigned to consumers in a way that facilitates join operations where required. (Range)
      • - * This can only be done if every consumer is subscribed to the same topics and the topics are co-partitioned. + * Properties are as follows: + *
        + *
      1. Each member must get at least one partition for every topic that it is subscribed to. The only exception is when + * the number of subscribed members is greater than the number of partitions for that topic. (Range)
        2. + *
      2. Partitions should be assigned to members in a way that facilitates the join operation when required. (Range)
        3. + * This can only be done if every member is subscribed to the same topics and the topics are co-partitioned. * Two streams are co-partitioned if the following conditions are met: - * ->The keys must have the same schemas - * ->The topics involved must have the same number of partitions - *
      3. 3) Consumers should retain as much as their previous assignment as possible. (Sticky)
        4. - *
    + *
      + *
    • The keys must have the same schemas. + *
    • The topics involved must have the same number of partitions. + *
    + *
  • Members should retain as much as their previous assignment as possible to reduce the number of partition movements during reassignment. (Sticky)
    • + * *

    * - *

    The algorithm works mainly in 5 steps described below - *

      - *
    • 1) Generate a map of consumersPerTopic using the member subscriptions.
      • - *
    • 2) Generate a list of consumers (potentiallyUnfilledConsumers) that have not met the minimum required quota for assignment AND + *

      The algorithm includes the following steps: + *

        + *
      1. Generate a map of membersPerTopic using the given member subscriptions.
        2. + *
      2. Generate a list of members (potentiallyUnfilledMembers) that have not met the minimum required quota for assignment AND * get a list of sticky partitions that we want to retain in the new assignment.
        3. - *
      3. 3) Add consumers from the potentiallyUnfilled list to the Unfilled list if they haven't met the total required quota i.e. minimum number of partitions per member + 1 (if member is designated to receive one of the excess partitions)
        4. - *
      4. 4) Generate a list of unassigned partitions by calculating the difference between total partitions and already assigned (sticky) partitions
        5. - *
      5. 5) Iterate through unfilled consumers and assign partitions from the available partitions
        6. - *
    + *
  • Add members from the potentiallyUnfilled list to the Unfilled list if they haven't met the total required quota i.e. minimum number of partitions per member + 1 (if member is designated to receive one of the excess partitions)
    • + *
  • Generate a list of unassigned partitions by calculating the difference between total partitions and already assigned (sticky) partitions
    • + *
  • Iterate through unfilled members and assign partitions from the unassigned partitions
    • + * *

    * */ @@ -65,119 +67,101 @@ public String name() { return RANGE_ASSIGNOR_NAME; } - private static void putList(Map> map, K key, V value) { - List list = map.computeIfAbsent(key, k -> new ArrayList<>()); - list.add(value); - } - - private static void putSet(Map> map, K key, V value) { - Set set = map.computeIfAbsent(key, k -> new HashSet<>()); - set.add(value); - } - - static class Pair { - private final T first; - private final U second; + static class RemainingAssignmentsForMember { + private final String memberId; + private final Integer remaining; - public Pair(T first, U second) { - this.first = first; - this.second = second; + public RemainingAssignmentsForMember(String memberId, Integer remaining) { + this.memberId = memberId; + this.remaining = remaining; } - public T getFirst() { - return first; + public String memberId() { + return memberId; } - public U getSecond() { - return second; + public Integer remaining() { + return remaining; } - @Override - public String toString() { - return "(" + first + ", " + second + ")"; - } } - // Returns a map of the list of consumers per Topic (keyed by topicId) - private Map> consumersPerTopic(AssignmentSpec assignmentSpec) { - Map> mapTopicsToConsumers = new HashMap<>(); - Map membersData = assignmentSpec.members; + private Map> membersPerTopic(final AssignmentSpec assignmentSpec) { + Map> membersPerTopic = new HashMap<>(); + Map membersData = assignmentSpec.members(); - for (Map.Entry memberEntry : membersData.entrySet()) { - String memberId = memberEntry.getKey(); - AssignmentMemberSpec memberMetadata = memberEntry.getValue(); - Collection topics = memberMetadata.subscribedTopics; + membersData.forEach((memberId, memberMetadata) -> { + Collection topics = memberMetadata.subscribedTopicIds(); for (Uuid topicId: topics) { - putList(mapTopicsToConsumers, topicId, memberId); + membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId); } - } - return mapTopicsToConsumers; + }); + + return membersPerTopic; } - private Map> getAvailablePartitionsPerTopic(AssignmentSpec assignmentSpec, Map> assignedStickyPartitionsPerTopic) { - Map> availablePartitionsPerTopic = new HashMap<>(); - Map topicsMetadata = assignmentSpec.topics; + private Map> getUnassignedPartitionsPerTopic(final AssignmentSpec assignmentSpec, Map> assignedStickyPartitionsPerTopic) { + Map> unassignedPartitionsPerTopic = new HashMap<>(); + Map topicsMetadata = assignmentSpec.topics(); - for (Map.Entry topicMetadataEntry : topicsMetadata.entrySet()) { - Uuid topicId = topicMetadataEntry.getKey(); - ArrayList availablePartitionsForTopic = new ArrayList<>(); - int numPartitions = topicsMetadata.get(topicId).numPartitions; - // since the loop iterates from 0 to n, the partitions will be in ascending order within the list of available partitions per topic + topicsMetadata.forEach((topicId, assignmentTopicMetadata) -> { + ArrayList unassignedPartitionsForTopic = new ArrayList<>(); + int numPartitions = assignmentTopicMetadata.numPartitions(); + // The partitions will be in ascending order within the list of unassigned partitions per topic. Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.getOrDefault(topicId, new HashSet<>()); for (int i = 0; i < numPartitions; i++) { if (!assignedStickyPartitionsForTopic.contains(i)) { - availablePartitionsForTopic.add(i); + unassignedPartitionsForTopic.add(i); } } - availablePartitionsPerTopic.put(topicId, availablePartitionsForTopic); - } - return availablePartitionsPerTopic; + unassignedPartitionsPerTopic.put(topicId, unassignedPartitionsForTopic); + }); + + return unassignedPartitionsPerTopic; } @Override - public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAssignorException { + public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws PartitionAssignorException { Map>> membersWithNewAssignmentPerTopic = new HashMap<>(); // Step 1 - Map> consumersPerTopic = consumersPerTopic(assignmentSpec); + Map> membersPerTopic = membersPerTopic(assignmentSpec); // Step 2 - Map>> unfilledConsumersPerTopic = new HashMap<>(); + Map>> unfilledMembersPerTopic = new HashMap<>(); Map> assignedStickyPartitionsPerTopic = new HashMap<>(); - for (Map.Entry> topicEntry : consumersPerTopic.entrySet()) { - Uuid topicId = topicEntry.getKey(); - // For each topic we have a temporary list of consumers stored in potentiallyUnfilledConsumers. - // The list is populated with consumers that satisfy one of the two conditions :- - // 1) Consumers that have the minimum required number of partitions .i.e numPartitionsPerConsumer BUT they could be assigned an extra partition later on. - // In this case we add the consumer to the unfilled consumers map iff an extra partition needs to be assigned to it. - // 2) Consumers that don't have the minimum required partitions, so irrespective of whether they get an extra partition or not they get added to the unfilled map later. - List> potentiallyUnfilledConsumers = new ArrayList<>(); - List consumersForTopic = topicEntry.getValue(); - - AssignmentTopicMetadata topicData = assignmentSpec.topics.get(topicId); - int numPartitionsForTopic = topicData.numPartitions; - // Initially, minRequiredQuota is the minimum number of partitions that each consumer should get i.e numPartitionsPerConsumer. - // Idle consumers case :- The numConsumers subscribed to a topic is greater than numPartitions. In such cases, all consumers get assigned via the "extra partitions" logic since min = 0. - int minRequiredQuota = numPartitionsForTopic / consumersForTopic.size(); - // Each consumer can get only one extra partition per topic after receiving the minimum quota = numPartitionsPerConsumer - int numConsumersWithExtraPartition = numPartitionsForTopic % consumersForTopic.size(); - - for (String memberId: consumersForTopic) { - - // Convert the set to a list first and sort the partitions in numeric order since we want the same partition numbers from each topic - // to go to the same consumer in case of co-partitioned topics. - Set currentAssignmentSetForTopic = assignmentSpec.members.get(memberId).currentAssignmentPerTopic.getOrDefault(topicId, new HashSet<>()); - // Size of the older assignment, this will be 0 when assign is called for the first time. - // The older assignment is required when a reassignment occurs to ensure stickiness. - int currentAssignmentSize = currentAssignmentSetForTopic.size(); - List currentAssignmentListForTopic = new ArrayList<>(currentAssignmentSetForTopic); - - // If there are previously assigned partitions present, we want to retain them. + membersPerTopic.forEach((topicId, membersForTopic) -> { + // For each topic we have a temporary list of members stored in potentiallyUnfilledMembers. + // The list is populated with members that satisfy one of the two conditions: + // 1) Members that already have the minimum required number of partitions i.e. total partitions divided by total subscribed members + // BUT they could be assigned an extra partition later on. Extra partitions exist if total partitions % number of members is greater than 0. + // In this case we add the member to the unfilled members map iff an extra partition needs to be assigned to it. + // 2) Members that don't have the minimum required partitions, so irrespective of whether they get an extra partition or not they get added to the unfilled map later. + List> potentiallyUnfilledMembers = new ArrayList<>(); + + AssignmentTopicMetadata topicData = assignmentSpec.topics().get(topicId); + int numPartitionsForTopic = topicData.numPartitions(); + + // Idle members case : When the number of members subscribed to a topic is greater than the total number of Partitions, + // all members get assigned via the "extra partitions" logic since minRequiredQuota = 0. + int minRequiredQuota = numPartitionsForTopic / membersForTopic.size(); + // Each member can get only ONE extra partition per topic after receiving the minimum quota. + int numMembersWithExtraPartition = numPartitionsForTopic % membersForTopic.size(); + + for (String memberId: membersForTopic) { + // The partitions need to be in numeric order since we want the same partition numbers from each topic + // to go to the same member in case of co-partitioned topics to facilitate joins. + Set assignedPartitionsForTopic = assignmentSpec.members().get(memberId).assignedPartitions().getOrDefault(topicId, new HashSet<>()); + + int currentAssignmentSize = assignedPartitionsForTopic.size(); + List currentAssignmentListForTopic = new ArrayList<>(assignedPartitionsForTopic); + + // If there were previously assigned partitions present, we want to retain them. if (currentAssignmentSize > 0) { // We either need to retain currentSize number of partitions when currentSize < required OR required number of partitions otherwise. int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); Collections.sort(currentAssignmentListForTopic); for (int i = 0; i < retainedPartitionsCount; i++) { - putSet(assignedStickyPartitionsPerTopic, topicId, currentAssignmentListForTopic.get(i)); + assignedStickyPartitionsPerTopic.computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); } } @@ -186,69 +170,60 @@ public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAss int remaining = minRequiredQuota - currentAssignmentSize; // There are 3 cases w.r.t value of remaining - // 1) remaining < 0 this means that the consumer has more than the min required amount. - // It could have an extra partition, so we check for that. - if (remaining < 0 && numConsumersWithExtraPartition > 0) { - // In order to remain as sticky as possible, since the order of members can be different, we want the consumers that already had extra - // partitions to retain them if it's still required, instead of assigning the extras to the first few consumers directly. - // Ex:- If two consumers out of 3 are supposed to get an extra partition and currently 1 of them already has the extra, we want this consumer - // to retain it first and later if we have partitions left they will be assigned to the first few consumers and the unfilled map is updated - numConsumersWithExtraPartition--; + // 1) remaining < 0 this means that the member has more than the min required amount. + if (remaining < 0 && numMembersWithExtraPartition > 0) { + // In order to remain as sticky as possible, since the order of members can be different, we want the members that already had extra + // partitions to retain them if it's still required, instead of assigning the extras to the first few members directly. + numMembersWithExtraPartition--; // Since we already added the minimumRequiredQuota of partitions in the previous step (until minReq - 1), we just need to - // add the extra partition which will be present at the index right after min quota is satisfied. - putSet(assignedStickyPartitionsPerTopic, topicId, currentAssignmentListForTopic.get(minRequiredQuota)); + // add the extra partition that will be present at the index right after min quota was satisfied. + assignedStickyPartitionsPerTopic.computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(minRequiredQuota)); membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(minRequiredQuota)); } else { - // 3) If remaining = 0 it has min req partitions but there is scope for getting an extra partition later on, so it is a potentialUnfilledConsumer. - // 4) If remaining > 0 it doesn't even have the min required partitions, and it definitely is unfilled so it should be added to potentialUnfilledConsumers. - Pair newPair = new Pair<>(memberId, remaining); - potentiallyUnfilledConsumers.add(newPair); + // 2) If remaining = 0 it has min req partitions but there is scope for getting an extra partition later on, so it is a potentialUnfilledMember. + // 3) If remaining > 0 it doesn't even have the min required partitions, and it definitely is unfilled, so it should be added to potentialUnfilledMembers. + RemainingAssignmentsForMember newPair = new RemainingAssignmentsForMember<>(memberId, remaining); + potentiallyUnfilledMembers.add(newPair); } } // Step 3 - // Iterate through potentially unfilled consumers and if remaining > 0 after considering the extra partitions assignment, add to the unfilled list. - for (Pair pair : potentiallyUnfilledConsumers) { - String memberId = pair.getFirst(); - Integer remaining = pair.getSecond(); - if (numConsumersWithExtraPartition > 0) { + // If remaining > 0 after increasing the required quota due to the extra partition, add potentially unfilled member to the unfilled members list. + for (RemainingAssignmentsForMember pair : potentiallyUnfilledMembers) { + String memberId = pair.memberId(); + Integer remaining = pair.remaining(); + if (numMembersWithExtraPartition > 0) { remaining++; - numConsumersWithExtraPartition--; + numMembersWithExtraPartition--; } if (remaining > 0) { - Pair newPair = new Pair<>(memberId, remaining); - putList(unfilledConsumersPerTopic, topicId, newPair); + RemainingAssignmentsForMember newPair = new RemainingAssignmentsForMember<>(memberId, remaining); + unfilledMembersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(newPair); } } - } + }); // Step 4 - // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get available partitions. - Map> availablePartitionsPerTopic = getAvailablePartitionsPerTopic(assignmentSpec, assignedStickyPartitionsPerTopic); + // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get unassigned partitions. + Map> unassignedPartitionsPerTopic = getUnassignedPartitionsPerTopic(assignmentSpec, assignedStickyPartitionsPerTopic); // Step 5 - // Iterate through the unfilled consumers list and assign remaining number of partitions from the availablePartitions list - for (Map.Entry>> unfilledEntry : unfilledConsumersPerTopic.entrySet()) { - Uuid topicId = unfilledEntry.getKey(); - List> unfilledConsumersForTopic = unfilledEntry.getValue(); - int newStartPointer = 0; - for (Pair stringIntegerPair : unfilledConsumersForTopic) { - String memberId = stringIntegerPair.getFirst(); - int remaining = stringIntegerPair.getSecond(); - // assign the first few partitions from the list and then delete them from the availablePartitions list for this topic - List subList = availablePartitionsPerTopic.get(topicId).subList(newStartPointer, newStartPointer + remaining); - newStartPointer += remaining; - membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).addAll(subList); + // Iterate through the unfilled members list and assign the remaining number of partitions from the unassignedPartitions list. + unfilledMembersPerTopic.forEach((topicId, unfilledMembersForTopic) -> { + int unassignedPartitionsListStartPointer = 0; + for (RemainingAssignmentsForMember remainingAssignmentsForMember : unfilledMembersForTopic) { + String memberId = remainingAssignmentsForMember.memberId(); + int remaining = remainingAssignmentsForMember.remaining(); + List partitionsToAssign = unassignedPartitionsPerTopic.get(topicId).subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); + unassignedPartitionsListStartPointer += remaining; + membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).addAll(partitionsToAssign); } - } + }); - // Consolidate the maps into MemberAssignment and then finally map each consumer to a MemberAssignment. + // Consolidate the maps into MemberAssignment and then finally map each member to a MemberAssignment. Map membersWithNewAssignment = new HashMap<>(); - for (Map.Entry>> consumer : membersWithNewAssignmentPerTopic.entrySet()) { - String consumerId = consumer.getKey(); - Map> assignmentPerTopic = consumer.getValue(); - membersWithNewAssignment.computeIfAbsent(consumerId, k -> new MemberAssignment(assignmentPerTopic)); - } + + membersWithNewAssignmentPerTopic.forEach((memberId, assignmentPerTopic) -> membersWithNewAssignment.computeIfAbsent(memberId, k -> new MemberAssignment(assignmentPerTopic))); return new GroupAssignment(membersWithNewAssignment); } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 38224ff03df26..2beb12f99e4ba 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -63,7 +63,7 @@ public void testOneConsumerNoTopic() { AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment groupAssignment = assignor.assign(assignmentSpec); - assertTrue(groupAssignment.getMembers().isEmpty()); + assertTrue(groupAssignment.members().isEmpty()); } @Test @@ -203,7 +203,7 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); // Consumer C shouldn't get any assignment, due to stickiness A, B retain their assignments - assertNull(computedAssignment.getMembers().get(consumerC)); + assertNull(computedAssignment.members().get(consumerC)); assertAssignment(expectedAssignment, computedAssignment); assertCoPartitionJoinProperty(computedAssignment); } @@ -366,8 +366,8 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme // We have a set of sets with the partitions that should be distributed amongst the consumers, if it exists then remove it from the set. private void assertAssignment(Map>> expectedAssignment, GroupAssignment computedGroupAssignment) { - for (MemberAssignment member : computedGroupAssignment.getMembers().values()) { - Map> computedAssignmentForMember = member.getAssignmentPerTopic(); + for (MemberAssignment member : computedGroupAssignment.members().values()) { + Map> computedAssignmentForMember = member.targetPartitions(); for (Map.Entry> assignmentForTopic : computedAssignmentForMember.entrySet()) { Uuid topicId = assignmentForTopic.getKey(); Set assignmentPartitionsSet = assignmentForTopic.getValue(); @@ -378,8 +378,8 @@ private void assertAssignment(Map>> expectedAssignment, G } private void assertCoPartitionJoinProperty(GroupAssignment groupAssignment) { - for (MemberAssignment member : groupAssignment.getMembers().values()) { - Map> computedAssignmentForMember = member.getAssignmentPerTopic(); + for (MemberAssignment member : groupAssignment.members().values()) { + Map> computedAssignmentForMember = member.targetPartitions(); Set compareSet = new HashSet<>(); for (Set partitionsForTopicSet : computedAssignmentForMember.values()) { if (compareSet.isEmpty()) { From 3d9264547cf9ba9afe3c27ce333bd7fe332216cc Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Mon, 17 Apr 2023 10:58:16 -0700 Subject: [PATCH 21/44] Interface changes incorporated, added stickiness test and non-existent topic test. --- .../group/assignor/AssignmentMemberSpec.java | 74 ++-- .../group/assignor/AssignmentSpec.java | 6 +- .../assignor/AssignmentTopicMetadata.java | 31 +- .../group/assignor/GroupAssignment.java | 13 +- .../group/assignor/MemberAssignment.java | 31 +- .../group/assignor/RangeAssignor.java | 11 +- .../ServerSideStickyRangeAssignor.java | 257 ------------ .../group/assignor/RangeAssignorTest.java | 78 ++-- .../ServerSideStickyRangeAssignorTest.java | 392 ------------------ 9 files changed, 140 insertions(+), 753 deletions(-) delete mode 100644 group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java delete mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java index c9bd525c1cd7d..143b73bd605fb 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java @@ -19,7 +19,6 @@ import org.apache.kafka.common.Uuid; import java.util.Collection; -import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -32,66 +31,93 @@ public class AssignmentMemberSpec { /** * The instance ID if provided. */ - final Optional instanceId; + private final Optional instanceId; /** * The rack ID if provided. */ - final Optional rackId; + private final Optional rackId; /** - * The topicIds of topics that the member is subscribed to. + * Topics Ids that the member is subscribed to. */ - final Collection subscribedTopics; + private final Collection subscribedTopicIds; /** - * Maps the partitions assigned for this member per topicId + * Partitions assigned keyed by topicId. */ - final Map> currentAssignmentPerTopic; + private final Map> assignedPartitions; + + /** + * @return The instance ID as an Optional. + */ + public Optional instanceId() { + return instanceId; + } + + /** + * @return The rack ID as an Optional. + */ + public Optional rackId() { + return rackId; + } + + /** + * @return Collection of subscribed topic Ids. + */ + public Collection subscribedTopicIds() { + return subscribedTopicIds; + } + + /** + * @return Assigned partitions keyed by topic Ids. + */ + public Map> assignedPartitions() { + return assignedPartitions; + } public AssignmentMemberSpec( - Optional instanceId, - Optional rackId, - List subscribedTopics, - Map> currentAssignmentPerTopic + Optional instanceId, + Optional rackId, + Collection subscribedTopicIds, + Map> assignedPartitions ) { Objects.requireNonNull(instanceId); Objects.requireNonNull(rackId); - Objects.requireNonNull(subscribedTopics); + Objects.requireNonNull(subscribedTopicIds); + Objects.requireNonNull(assignedPartitions); this.instanceId = instanceId; this.rackId = rackId; - this.subscribedTopics = subscribedTopics; - this.currentAssignmentPerTopic = currentAssignmentPerTopic; + this.subscribedTopicIds = subscribedTopicIds; + this.assignedPartitions = assignedPartitions; } @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; - AssignmentMemberSpec that = (AssignmentMemberSpec) o; - if (!instanceId.equals(that.instanceId)) return false; if (!rackId.equals(that.rackId)) return false; - if (!subscribedTopics.equals(that.subscribedTopics)) return false; - return currentAssignmentPerTopic.equals(that.currentAssignmentPerTopic); + if (!subscribedTopicIds.equals(that.subscribedTopicIds)) return false; + return assignedPartitions.equals(that.assignedPartitions); } @Override public int hashCode() { int result = instanceId.hashCode(); result = 31 * result + rackId.hashCode(); - result = 31 * result + subscribedTopics.hashCode(); - result = 31 * result + currentAssignmentPerTopic.hashCode(); + result = 31 * result + subscribedTopicIds.hashCode(); + result = 31 * result + assignedPartitions.hashCode(); return result; } @Override public String toString() { return "AssignmentMemberSpec(instanceId=" + instanceId + - ", rackId=" + rackId + - ", subscribedTopics=" + subscribedTopics + - ", currentAssignment=" + currentAssignmentPerTopic + - ')'; + ", rackId=" + rackId + + ", subscribedTopicIds=" + subscribedTopicIds + + ", assignedPartitions=" + assignedPartitions + + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java index 39426dadfcce1..03fd663a40465 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java @@ -63,9 +63,7 @@ public Map topics() { public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; - AssignmentSpec that = (AssignmentSpec) o; - if (!members.equals(that.members)) return false; return topics.equals(that.topics); } @@ -80,7 +78,7 @@ public int hashCode() { @Override public String toString() { return "AssignmentSpec(members=" + members + - ", topics=" + topics + - ')'; + ", topics=" + topics + + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java index 4fb9c63c02eaa..283ad7f8bdec4 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentTopicMetadata.java @@ -16,16 +16,10 @@ */ package org.apache.kafka.coordinator.group.assignor; -import java.util.Objects; - /** * Metadata of a topic. */ public class AssignmentTopicMetadata { - /** - * The topic name. - */ - private final String topicName; /** * The number of partitions. @@ -33,21 +27,11 @@ public class AssignmentTopicMetadata { private final int numPartitions; public AssignmentTopicMetadata( - String topicName, - int numPartitions + int numPartitions ) { - Objects.requireNonNull(topicName); - this.topicName = topicName; this.numPartitions = numPartitions; } - /** - * @return The topic name. - */ - public String topicName() { - return topicName; - } - /** * @return The number of partitions present for the topic. */ @@ -59,24 +43,17 @@ public int numPartitions() { public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; - AssignmentTopicMetadata that = (AssignmentTopicMetadata) o; - - if (numPartitions != that.numPartitions) return false; - return topicName.equals(that.topicName); + return numPartitions == that.numPartitions; } @Override public int hashCode() { - int result = topicName.hashCode(); - result = 31 * result + numPartitions; - return result; + return numPartitions; } @Override public String toString() { - return "AssignmentTopicMetadata(topicName=" + topicName + - ", numPartitions=" + numPartitions + - ')'; + return "AssignmentTopicMetadata(numPartitions=" + numPartitions + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java index 6256b7453a835..f6ba5cac20821 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java @@ -26,26 +26,27 @@ public class GroupAssignment { /** * The member assignments keyed by member id. */ - final Map members; + private final Map members; public GroupAssignment( - Map members + Map members ) { Objects.requireNonNull(members); this.members = members; } - public Map getMembers() { - return this.members; + /** + * @return Member assignments keyed by member Ids. + */ + public Map members() { + return members; } @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; - GroupAssignment that = (GroupAssignment) o; - return members.equals(that.members); } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java index 3f9dd7f8d4378..ca19b3e883b70 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/MemberAssignment.java @@ -22,41 +22,42 @@ import java.util.Objects; import java.util.Set; - /** * The partition assignment for a consumer group member. */ public class MemberAssignment { - - private final Map> assignmentPerTopic; - - - public MemberAssignment(Map> assignmentPerTopic) { - Objects.requireNonNull(assignmentPerTopic); - this.assignmentPerTopic = assignmentPerTopic; + /** + * The target partitions assigned to this member keyed by topicId. + */ + private final Map> targetPartitions; + + public MemberAssignment(Map> targetPartitions) { + Objects.requireNonNull(targetPartitions); + this.targetPartitions = targetPartitions; } - public Map> getAssignmentPerTopic() { - return this.assignmentPerTopic; + /** + * @return Target partitions keyed by topic Ids. + */ + public Map> targetPartitions() { + return this.targetPartitions; } @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; - MemberAssignment that = (MemberAssignment) o; - - return assignmentPerTopic.equals(that.assignmentPerTopic); + return targetPartitions.equals(that.targetPartitions); } @Override public int hashCode() { - return assignmentPerTopic.hashCode(); + return targetPartitions.hashCode(); } @Override public String toString() { - return "MemberAssignment(targetPartitions=" + assignmentPerTopic + ')'; + return "MemberAssignment(targetPartitions=" + targetPartitions + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index b44556ae1fce6..de0255d7c28e1 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -17,6 +17,8 @@ package org.apache.kafka.coordinator.group.assignor; import org.apache.kafka.common.Uuid; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.util.ArrayList; import java.util.Collection; @@ -60,6 +62,8 @@ */ public class RangeAssignor implements PartitionAssignor { + private static final Logger log = LoggerFactory.getLogger(RangeAssignor.class); + public static final String RANGE_ASSIGNOR_NAME = "range"; @Override @@ -93,7 +97,12 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS membersData.forEach((memberId, memberMetadata) -> { Collection topics = memberMetadata.subscribedTopicIds(); for (Uuid topicId: topics) { - membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId); + // Only topics that are present in both the subscribed topics list and the topic metadata should be considered for assignment. + if (assignmentSpec.topics().containsKey(topicId)) { + membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId); + } else { + log.info(memberId + " subscribed to topic " + topicId + " which doesn't exist in the topic metadata"); + } } }); diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java deleted file mode 100644 index bb89296d0c0ee..0000000000000 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignor.java +++ /dev/null @@ -1,257 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.kafka.coordinator.group.assignor; - -import org.apache.kafka.common.Uuid; - -import java.util.ArrayList; -import java.util.Collection; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; - -import static java.lang.Math.min; - -/** - *

    The Server Side Sticky Range Assignor inherits properties of both the range assignor and the sticky assignor. - * Properties are :- - *

      - *
    • 1) Each consumer must get at least one partition per topic that it is subscribed to whenever the number of consumers is - * less than or equal to the number of partitions for that topic. (Range)
      • - *
    • 2) Partitions should be assigned to consumers in a way that facilitates join operations where required. (Range)
      • - * This can only be done if the topics are co-partitioned in the first place - * Co-partitioned:- - * Two streams are co-partitioned if the following conditions are met:- - * ->The keys must have the same schemas - * ->The topics involved must have the same number of partitions - *
    • 3) Consumers should retain as much as their previous assignment as possible. (Sticky)
      • - *
    - *

    - * - *

    The algorithm works mainly in 5 steps described below - *

      - *
    • 1) Get a map of the consumersPerTopic created using the member subscriptions.
      • - *
    • 2) Get a list of consumers (potentiallyUnfilled) that have not met the minimum required quota for assignment AND - * get a list of sticky partitions that we want to retain in the new assignment.
      • - *
    • 3) Add consumers from potentiallyUnfilled to Unfilled if they haven't met the total required quota = minQuota + (if necessary) extraPartition
      • - *
    • 4) Get a list of available partitions by calculating the difference between total partitions and assigned sticky partitions
      • - *
    • 5) Iterate through unfilled consumers and assign partitions from available partitions
      • - *
    - *

    - * - */ -public class ServerSideStickyRangeAssignor implements PartitionAssignor { - - public static final String RANGE_ASSIGNOR_NAME = "range"; - - @Override - public String name() { - return RANGE_ASSIGNOR_NAME; - } - - private static void putList(Map> map, K key, V value) { - List list = map.computeIfAbsent(key, k -> new ArrayList<>()); - list.add(value); - } - - private static void putSet(Map> map, K key, V value) { - Set set = map.computeIfAbsent(key, k -> new HashSet<>()); - set.add(value); - } - - static class Pair { - private final T first; - private final U second; - - public Pair(T first, U second) { - this.first = first; - this.second = second; - } - - public T getFirst() { - return first; - } - - public U getSecond() { - return second; - } - - @Override - public String toString() { - return "(" + first + ", " + second + ")"; - } - } - - // Returns a map of the list of consumers per Topic (keyed by topicId) - private Map> consumersPerTopic(AssignmentSpec assignmentSpec) { - Map> mapTopicsToConsumers = new HashMap<>(); - Map membersData = assignmentSpec.members; - - for (Map.Entry memberEntry : membersData.entrySet()) { - String memberId = memberEntry.getKey(); - AssignmentMemberSpec memberMetadata = memberEntry.getValue(); - Collection topics = memberMetadata.subscribedTopics; - for (Uuid topicId: topics) { - putList(mapTopicsToConsumers, topicId, memberId); - } - } - return mapTopicsToConsumers; - } - - private Map> getAvailablePartitionsPerTopic(AssignmentSpec assignmentSpec, Map> assignedStickyPartitionsPerTopic) { - Map> availablePartitionsPerTopic = new HashMap<>(); - Map topicsMetadata = assignmentSpec.topics; - - for (Map.Entry topicMetadataEntry : topicsMetadata.entrySet()) { - Uuid topicId = topicMetadataEntry.getKey(); - ArrayList availablePartitionsForTopic = new ArrayList<>(); - int numPartitions = topicsMetadata.get(topicId).numPartitions; - // since the loop iterates from 0 to n, the partitions will be in ascending order within the list of available partitions per topic - Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.getOrDefault(topicId, new HashSet<>()); - for (int i = 0; i < numPartitions; i++) { - if (!assignedStickyPartitionsForTopic.contains(i)) { - availablePartitionsForTopic.add(i); - } - } - availablePartitionsPerTopic.put(topicId, availablePartitionsForTopic); - } - return availablePartitionsPerTopic; - } - - @Override - public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAssignorException { - Map>> membersWithNewAssignmentPerTopic = new HashMap<>(); - // Step 1 - Map> consumersPerTopic = consumersPerTopic(assignmentSpec); - // Step 2 - Map>> unfilledConsumersPerTopic = new HashMap<>(); - Map> assignedStickyPartitionsPerTopic = new HashMap<>(); - - for (Map.Entry> topicEntry : consumersPerTopic.entrySet()) { - Uuid topicId = topicEntry.getKey(); - // For each topic we have a temporary list of consumers stored in potentiallyUnfilledConsumers. - // The list is populated with consumers that satisfy one of the two conditions :- - // 1) Consumers that have the minimum required number of partitions .i.e numPartitionsPerConsumer BUT they could be assigned an extra partition later on. - // In this case we add the consumer to the unfilled consumers map iff an extra partition needs to be assigned to it. - // 2) Consumers that don't have the minimum required partitions, so irrespective of whether they get an extra partition or not they get added to the unfilled map later. - List> potentiallyUnfilledConsumers = new ArrayList<>(); - List consumersForTopic = topicEntry.getValue(); - - AssignmentTopicMetadata topicData = assignmentSpec.topics.get(topicId); - int numPartitionsForTopic = topicData.numPartitions; - // Initially, minRequiredQuota is the minimum number of partitions that each consumer should get i.e numPartitionsPerConsumer. - // Idle consumers case :- The numConsumers subscribed to a topic is greater than numPartitions. In such cases, all consumers get assigned via the "extra partitions" logic since min = 0. - int minRequiredQuota = numPartitionsForTopic / consumersForTopic.size(); - // Each consumer can get only one extra partition per topic after receiving the minimum quota = numPartitionsPerConsumer - int numConsumersWithExtraPartition = numPartitionsForTopic % consumersForTopic.size(); - - for (String memberId: consumersForTopic) { - - // Convert the set to a list first and sort the partitions in numeric order since we want the same partition numbers from each topic - // to go to the same consumer in case of co-partitioned topics. - Set currentAssignmentSetForTopic = assignmentSpec.members.get(memberId).currentAssignmentPerTopic.getOrDefault(topicId, new HashSet<>()); - // Size of the older assignment, this will be 0 when assign is called for the first time. - // The older assignment is required when a reassignment occurs to ensure stickiness. - int currentAssignmentSize = currentAssignmentSetForTopic.size(); - List currentAssignmentListForTopic = new ArrayList<>(currentAssignmentSetForTopic); - - // If there are previously assigned partitions present, we want to retain them. - if (currentAssignmentSize > 0) { - // We either need to retain currentSize number of partitions when currentSize < required OR required number of partitions otherwise. - int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); - Collections.sort(currentAssignmentListForTopic); - for (int i = 0; i < retainedPartitionsCount; i++) { - putSet(assignedStickyPartitionsPerTopic, topicId, currentAssignmentListForTopic.get(i)); - membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); - } - } - - // Number of partitions left to reach the minRequiredQuota. - int remaining = minRequiredQuota - currentAssignmentSize; - - // There are 3 cases w.r.t value of remaining - // 1) remaining < 0 this means that the consumer has more than the min required amount. - // It could have an extra partition, so we check for that. - if (remaining < 0 && numConsumersWithExtraPartition > 0) { - // In order to remain as sticky as possible, since the order of members can be different, we want the consumers that already had extra - // partitions to retain them if it's still required, instead of assigning the extras to the first few consumers directly. - // Ex:- If two consumers out of 3 are supposed to get an extra partition and currently 1 of them already has the extra, we want this consumer - // to retain it first and later if we have partitions left they will be assigned to the first few consumers and the unfilled map is updated - numConsumersWithExtraPartition--; - // Since we already added the minimumRequiredQuota of partitions in the previous step (until minReq - 1), we just need to - // add the extra partition which will be present at the index right after min quota is satisfied. - putSet(assignedStickyPartitionsPerTopic, topicId, currentAssignmentListForTopic.get(minRequiredQuota)); - membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(minRequiredQuota)); - } else { - // 3) If remaining = 0 it has min req partitions but there is scope for getting an extra partition later on, so it is a potentialUnfilledConsumer. - // 4) If remaining > 0 it doesn't even have the min required partitions, and it definitely is unfilled so it should be added to potentialUnfilledConsumers. - Pair newPair = new Pair<>(memberId, remaining); - potentiallyUnfilledConsumers.add(newPair); - } - } - - // Step 3 - // Iterate through potentially unfilled consumers and if remaining > 0 after considering the extra partitions assignment, add to the unfilled list. - for (Pair pair : potentiallyUnfilledConsumers) { - String memberId = pair.getFirst(); - Integer remaining = pair.getSecond(); - if (numConsumersWithExtraPartition > 0) { - remaining++; - numConsumersWithExtraPartition--; - } - if (remaining > 0) { - Pair newPair = new Pair<>(memberId, remaining); - putList(unfilledConsumersPerTopic, topicId, newPair); - } - } - } - - // Step 4 - // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get available partitions. - Map> availablePartitionsPerTopic = getAvailablePartitionsPerTopic(assignmentSpec, assignedStickyPartitionsPerTopic); - - // Step 5 - // Iterate through the unfilled consumers list and assign remaining number of partitions from the availablePartitions list - for (Map.Entry>> unfilledEntry : unfilledConsumersPerTopic.entrySet()) { - Uuid topicId = unfilledEntry.getKey(); - List> unfilledConsumersForTopic = unfilledEntry.getValue(); - int newStartPointer = 0; - for (Pair stringIntegerPair : unfilledConsumersForTopic) { - String memberId = stringIntegerPair.getFirst(); - int remaining = stringIntegerPair.getSecond(); - // assign the first few partitions from the list and then delete them from the availablePartitions list for this topic - List subList = availablePartitionsPerTopic.get(topicId).subList(newStartPointer, newStartPointer + remaining); - newStartPointer += remaining; - membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).addAll(subList); - } - } - - // Consolidate the maps into MemberAssignment and then finally map each consumer to a MemberAssignment. - Map membersWithNewAssignment = new HashMap<>(); - for (Map.Entry>> consumer : membersWithNewAssignmentPerTopic.entrySet()) { - String consumerId = consumer.getKey(); - Map> assignmentPerTopic = consumer.getValue(); - membersWithNewAssignment.computeIfAbsent(consumerId, k -> new MemberAssignment(assignmentPerTopic)); - } - - return new GroupAssignment(membersWithNewAssignment); - } -} - diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 2beb12f99e4ba..77f9628ba0b29 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -36,18 +36,10 @@ import static org.junit.jupiter.api.Assertions.assertTrue; public class RangeAssignorTest { - private final RangeAssignor assignor = new RangeAssignor(); - - private final String topic1Name = "topic1"; private final Uuid topic1Uuid = Uuid.randomUuid(); - - private final String topic2Name = "topic2"; private final Uuid topic2Uuid = Uuid.randomUuid(); - - private final String topic3Name = "topic3"; private final Uuid topic3Uuid = Uuid.randomUuid(); - private final String consumerA = "A"; private final String consumerB = "B"; private final String consumerC = "C"; @@ -55,7 +47,7 @@ public class RangeAssignorTest { @Test public void testOneConsumerNoTopic() { Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); Map members = new HashMap<>(); List subscribedTopics = new ArrayList<>(); members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopics, new HashMap<>())); @@ -66,13 +58,28 @@ public void testOneConsumerNoTopic() { assertTrue(groupAssignment.members().isEmpty()); } + @Test + public void testOneConsumerNonExistentTopic() { + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); + Map members = new HashMap<>(); + List subscribedTopics = new ArrayList<>(); + subscribedTopics.add(topic2Uuid); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopics, new HashMap<>())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment groupAssignment = assignor.assign(assignmentSpec); + + assertTrue(groupAssignment.members().isEmpty()); + } + @Test public void testFirstAssignmentTwoConsumersTwoTopicsSameSubscriptions() { // A -> T1, T3 // B -> T1, T3 // T1 -> 3 Partitions // T3 -> 2 Partitions // Topics Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); + topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); // Members Map members = new HashMap<>(); // Consumer A @@ -101,9 +108,9 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() // A -> T1, T2 // B -> T3 // C -> T2, T3 // T1 -> 3 Partitions // T2 -> 3 Partitions // T3 -> 2 Partitions // Topics Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); - topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); + topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); // Members Map members = new HashMap<>(); // Consumer A @@ -136,9 +143,9 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { // Topics Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); + topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); // Topic 3 has 2 partitions but three consumers subscribed to it - one of them will not get an assignment - topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); // Members Map members = new HashMap<>(); // Consumer A @@ -170,8 +177,8 @@ public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerAdded() { // Topics Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 2)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 2)); + topics.put(topic1Uuid, new AssignmentTopicMetadata(2)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(2)); // Members Map members = new HashMap<>(); // Add a new consumer to trigger a re-assignment @@ -202,6 +209,10 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + // Test for stickiness + assertEquals(computedAssignment.members().get(consumerA).targetPartitions(), new HashMap<>(currentAssignmentForA), "Stickiness test failed for Consumer A"); + assertEquals(computedAssignment.members().get(consumerB).targetPartitions(), new HashMap<>(currentAssignmentForB), "Stickiness test failed for Consumer B"); + // Consumer C shouldn't get any assignment, due to stickiness A, B retain their assignments assertNull(computedAssignment.members().get(consumerC)); assertAssignment(expectedAssignment, computedAssignment); @@ -228,8 +239,8 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { // Simulating adding a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 4)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 4)); + topics.put(topic1Uuid, new AssignmentTopicMetadata(4)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(4)); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); @@ -242,6 +253,10 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(2, 3))); + // Test for stickiness + assertEquals(computedAssignment.members().get(consumerA).targetPartitions(), new HashMap<>(currentAssignmentForA), "Stickiness test failed for Consumer A"); + // Implicitly B also retained its partitions since the other set of assigned partitions is (2,3) + assertAssignment(expectedAssignment, computedAssignment); assertCoPartitionJoinProperty(computedAssignment); } @@ -249,8 +264,8 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { @Test public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoConsumersTwoTopics() { Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); + topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); Map members = new HashMap<>(); // Consumer A @@ -283,6 +298,11 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + // Test for stickiness + assertEquals(computedAssignment.members().get(consumerB).targetPartitions(), new HashMap<>(currentAssignmentForB), "Stickiness test failed for Consumer B"); + assertTrue(computedAssignment.members().get(consumerA).targetPartitions().get(topic1Uuid).contains(0), "Stickiness test failed for Consumer A"); + // Co-partition join property test ensures that A retained the same partition for all topics. + assertAssignment(expectedAssignment, computedAssignment); assertCoPartitionJoinProperty(computedAssignment); } @@ -290,8 +310,8 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon @Test public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoConsumersTwoTopics() { Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); + topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); Map members = new HashMap<>(); // Consumer A was removed @@ -320,9 +340,9 @@ public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoC public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignmentWithThreeConsumersTwoTopics() { Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); - topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); + topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); + topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3 // Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2 @@ -361,6 +381,10 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme // Topic 3 Partitions Assignment expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + // Test for stickiness + assertTrue(computedAssignment.members().get(consumerC).targetPartitions().get(topic2Uuid).contains(2), "Stickiness test failed for Consumer C"); + assertTrue(computedAssignment.members().get(consumerA).targetPartitions().get(topic1Uuid).containsAll(Arrays.asList(0, 1)), "Stickiness test failed for Consumer A"); + assertAssignment(expectedAssignment, computedAssignment); } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java deleted file mode 100644 index 52d36eece5610..0000000000000 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/ServerSideStickyRangeAssignorTest.java +++ /dev/null @@ -1,392 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.kafka.coordinator.group.assignor; - -import org.apache.kafka.common.Uuid; - -import org.junit.jupiter.api.Test; - -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.Set; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertTrue; - -public class ServerSideStickyRangeAssignorTest { - - private final ServerSideStickyRangeAssignor assignor = new ServerSideStickyRangeAssignor(); - - private final String topic1Name = "topic1"; - private final Uuid topic1Uuid = Uuid.randomUuid(); - - private final String topic2Name = "topic2"; - private final Uuid topic2Uuid = Uuid.randomUuid(); - - private final String topic3Name = "topic3"; - private final Uuid topic3Uuid = Uuid.randomUuid(); - - private final String consumerA = "A"; - private final String consumerB = "B"; - private final String consumerC = "C"; - - @Test - public void testOneConsumerNoTopic() { - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - Map members = new HashMap<>(); - List subscribedTopics = new ArrayList<>(); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopics, new HashMap<>())); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment groupAssignment = assignor.assign(assignmentSpec); - - assertTrue(groupAssignment.getMembers().isEmpty()); - } - - @Test - public void testFirstAssignmentTwoConsumersTwoTopicsSameSubscriptions() { - // A -> T1, T3 // B -> T1, T3 // T1 -> 3 Partitions // T3 -> 2 Partitions - // Topics - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); - // Members - Map members = new HashMap<>(); - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); - members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); - members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - - Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); - // Topic 3 Partitions Assignment - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - - assertAssignment(expectedAssignment, computedAssignment); - } - - @Test - public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() { - // A -> T1, T2 // B -> T3 // C -> T2, T3 // T1 -> 3 Partitions // T2 -> 3 Partitions // T3 -> 2 Partitions - // Topics - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); - topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); - // Members - Map members = new HashMap<>(); - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); - // Consumer B - List subscribedTopicsB = new ArrayList<>(Collections.singletonList(topic3Uuid)); - members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); - // Consumer C - List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic2Uuid, topic3Uuid)); - members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - - Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); - // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); - // Topic 3 Partitions Assignment - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - - assertAssignment(expectedAssignment, computedAssignment); - } - - @Test - public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { - // Topics - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - // Topic 3 has 2 partitions but three consumers subscribed to it - one of them will not get an assignment - topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); - // Members - Map members = new HashMap<>(); - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); - members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); - members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); - // Consumer C - List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); - members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - - Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); - // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - - assertAssignment(expectedAssignment, computedAssignment); - } - - @Test - public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerAdded() { - // Topics - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 2)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 2)); - // Members - Map members = new HashMap<>(); - // Add a new consumer to trigger a re-assignment - // Consumer C - List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Collections.singletonList(0))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Collections.singletonList(0))); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(1))); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(1))); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - - Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - - // Consumer C shouldn't get any assignment, due to stickiness A, B retain their assignments - assertNull(computedAssignment.getMembers().get(consumerC)); - assertAssignment(expectedAssignment, computedAssignment); - assertCoPartitionJoinProperty(computedAssignment); - } - - @Test - public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { - // T1, T2 both have 3 partitions each when first assignment was calculated -> currentAssignmentForX - Map members = new HashMap<>(); - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); - - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); - - // Simulating adding a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 4)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 4)); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - - Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(2, 3))); - // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(2, 3))); - - assertAssignment(expectedAssignment, computedAssignment); - assertCoPartitionJoinProperty(computedAssignment); - } - - @Test - public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoConsumersTwoTopics() { - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); - - Map members = new HashMap<>(); - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); - // Add a new consumer to trigger a re-assignment - // Consumer C - List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - - Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - - assertAssignment(expectedAssignment, computedAssignment); - assertCoPartitionJoinProperty(computedAssignment); - } - - @Test - public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoConsumersTwoTopics() { - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); - - Map members = new HashMap<>(); - // Consumer A was removed - - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - - Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); - // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); - - assertAssignment(expectedAssignment, computedAssignment); - assertCoPartitionJoinProperty(computedAssignment); - } - - @Test - public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignmentWithThreeConsumersTwoTopics() { - - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(topic1Name, 3)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(topic2Name, 3)); - topics.put(topic3Uuid, new AssignmentTopicMetadata(topic3Name, 2)); - - // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3 - // Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2 - Map members = new HashMap<>(); - // Consumer A - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1, 2))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Collections.singletonList(0))); - // Change subscriptions - List subscribedTopicsA = new ArrayList<>(Collections.singletonList(topic1Uuid)); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); - // Consumer B - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(1))); - // Change subscriptions - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid, topic3Uuid)); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); - // Consumer C - Map> currentAssignmentForC = new HashMap<>(); - currentAssignmentForC.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); - currentAssignmentForC.put(topic3Uuid, new HashSet<>(Arrays.asList(0, 1))); - // Change subscriptions - List subscribedTopicsC = new ArrayList<>(Collections.singletonList(topic2Uuid)); - members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); - - AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - - Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); - // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); - // Topic 3 Partitions Assignment - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - - assertAssignment(expectedAssignment, computedAssignment); - } - - // We have a set of sets with the partitions that should be distributed amongst the consumers, if it exists then remove it from the set. - private void assertAssignment(Map>> expectedAssignment, GroupAssignment computedGroupAssignment) { - for (MemberAssignment member : computedGroupAssignment.getMembers().values()) { - Map> computedAssignmentForMember = member.getAssignmentPerTopic(); - for (Map.Entry> assignmentForTopic : computedAssignmentForMember.entrySet()) { - Uuid topicId = assignmentForTopic.getKey(); - Set assignmentPartitionsSet = assignmentForTopic.getValue(); - assertTrue(expectedAssignment.get(topicId).contains(assignmentPartitionsSet)); - expectedAssignment.remove(assignmentPartitionsSet); - } - } - } - - private void assertCoPartitionJoinProperty(GroupAssignment groupAssignment) { - for (MemberAssignment member : groupAssignment.getMembers().values()) { - Map> computedAssignmentForMember = member.getAssignmentPerTopic(); - Set compareSet = new HashSet<>(); - for (Set partitionsForTopicSet : computedAssignmentForMember.values()) { - if (compareSet.isEmpty()) { - compareSet = partitionsForTopicSet; - } - assertEquals(compareSet, partitionsForTopicSet); - } - } - } -} From 27a86d8bfc23d86a50067df1749bb1247351577b Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Wed, 19 Apr 2023 11:10:32 -0700 Subject: [PATCH 22/44] addressed some PR comments --- .../group/assignor/AssignmentMemberSpec.java | 1 - .../group/assignor/GroupAssignment.java | 2 +- .../group/assignor/RangeAssignor.java | 56 ++++++++++++------- .../group/assignor/RangeAssignorTest.java | 19 +++---- 4 files changed, 47 insertions(+), 31 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java index f8508bfc09f25..6de0978c41782 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java @@ -119,6 +119,5 @@ public String toString() { ", subscribedTopicIds=" + subscribedTopicIds + ", assignedPartitions=" + assignedPartitions + ')'; - } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java index f6ba5cac20821..f5cee5bf2e921 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GroupAssignment.java @@ -29,7 +29,7 @@ public class GroupAssignment { private final Map members; public GroupAssignment( - Map members + Map members ) { Objects.requireNonNull(members); this.members = members; diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index de0255d7c28e1..8ed67361ef86d 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -48,17 +48,6 @@ * *

    * - *

    The algorithm includes the following steps: - *

      - *
    1. Generate a map of membersPerTopic using the given member subscriptions.
      2. - *
    2. Generate a list of members (potentiallyUnfilledMembers) that have not met the minimum required quota for assignment AND - * get a list of sticky partitions that we want to retain in the new assignment.
      3. - *
    3. Add members from the potentiallyUnfilled list to the Unfilled list if they haven't met the total required quota i.e. minimum number of partitions per member + 1 (if member is designated to receive one of the excess partitions)
      4. - *
    4. Generate a list of unassigned partitions by calculating the difference between total partitions and already assigned (sticky) partitions
      5. - *
    5. Iterate through unfilled members and assign partitions from the unassigned partitions
      6. - *
    - *

    - * */ public class RangeAssignor implements PartitionAssignor { @@ -99,7 +88,9 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS for (Uuid topicId: topics) { // Only topics that are present in both the subscribed topics list and the topic metadata should be considered for assignment. if (assignmentSpec.topics().containsKey(topicId)) { - membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId); + membersPerTopic + .computeIfAbsent(topicId, k -> new ArrayList<>()) + .add(memberId); } else { log.info(memberId + " subscribed to topic " + topicId + " which doesn't exist in the topic metadata"); } @@ -129,6 +120,18 @@ private Map> getUnassignedPartitionsPerTopic(final Assignmen return unassignedPartitionsPerTopic; } +/** + *

    The algorithm includes the following steps: + *

      + *
    1. Generate a map of membersPerTopic using the given member subscriptions.
      2. + *
    2. Generate a list of members (potentiallyUnfilledMembers) that have not met the minimum required quota for assignment AND + * get a list of sticky partitions that we want to retain in the new assignment.
      3. + *
    3. Add members from the potentiallyUnfilled list to the Unfilled list if they haven't met the total required quota i.e. minimum number of partitions per member + 1 (if member is designated to receive one of the excess partitions)
      4. + *
    4. Generate a list of unassigned partitions by calculating the difference between total partitions and already assigned (sticky) partitions
      5. + *
    5. Iterate through unfilled members and assign partitions from the unassigned partitions
      6. + *
    + *

    + */ @Override public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws PartitionAssignorException { Map>> membersWithNewAssignmentPerTopic = new HashMap<>(); @@ -159,7 +162,7 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit for (String memberId: membersForTopic) { // The partitions need to be in numeric order since we want the same partition numbers from each topic // to go to the same member in case of co-partitioned topics to facilitate joins. - Set assignedPartitionsForTopic = assignmentSpec.members().get(memberId).assignedPartitions().getOrDefault(topicId, new HashSet<>()); + Set assignedPartitionsForTopic = assignmentSpec.members().get(memberId).assignedPartitions().getOrDefault(topicId, Collections.emptySet()); int currentAssignmentSize = assignedPartitionsForTopic.size(); List currentAssignmentListForTopic = new ArrayList<>(assignedPartitionsForTopic); @@ -170,8 +173,13 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); Collections.sort(currentAssignmentListForTopic); for (int i = 0; i < retainedPartitionsCount; i++) { - assignedStickyPartitionsPerTopic.computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); - membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(i)); + assignedStickyPartitionsPerTopic + .computeIfAbsent(topicId, k -> new HashSet<>()) + .add(currentAssignmentListForTopic.get(i)); + membersWithNewAssignmentPerTopic + .computeIfAbsent(memberId, k -> new HashMap<>()) + .computeIfAbsent(topicId, k -> new HashSet<>()) + .add(currentAssignmentListForTopic.get(i)); } } @@ -186,8 +194,13 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit numMembersWithExtraPartition--; // Since we already added the minimumRequiredQuota of partitions in the previous step (until minReq - 1), we just need to // add the extra partition that will be present at the index right after min quota was satisfied. - assignedStickyPartitionsPerTopic.computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(minRequiredQuota)); - membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).add(currentAssignmentListForTopic.get(minRequiredQuota)); + assignedStickyPartitionsPerTopic + .computeIfAbsent(topicId, k -> new HashSet<>()) + .add(currentAssignmentListForTopic.get(minRequiredQuota)); + membersWithNewAssignmentPerTopic + .computeIfAbsent(memberId, k -> new HashMap<>()) + .computeIfAbsent(topicId, k -> new HashSet<>()) + .add(currentAssignmentListForTopic.get(minRequiredQuota)); } else { // 2) If remaining = 0 it has min req partitions but there is scope for getting an extra partition later on, so it is a potentialUnfilledMember. // 3) If remaining > 0 it doesn't even have the min required partitions, and it definitely is unfilled, so it should be added to potentialUnfilledMembers. @@ -207,7 +220,9 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit } if (remaining > 0) { RemainingAssignmentsForMember newPair = new RemainingAssignmentsForMember<>(memberId, remaining); - unfilledMembersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(newPair); + unfilledMembersPerTopic + .computeIfAbsent(topicId, k -> new ArrayList<>()) + .add(newPair); } } }); @@ -225,7 +240,10 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit int remaining = remainingAssignmentsForMember.remaining(); List partitionsToAssign = unassignedPartitionsPerTopic.get(topicId).subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); unassignedPartitionsListStartPointer += remaining; - membersWithNewAssignmentPerTopic.computeIfAbsent(memberId, k -> new HashMap<>()).computeIfAbsent(topicId, k -> new HashSet<>()).addAll(partitionsToAssign); + membersWithNewAssignmentPerTopic + .computeIfAbsent(memberId, k -> new HashMap<>()) + .computeIfAbsent(topicId, k -> new HashSet<>()) + .addAll(partitionsToAssign); } }); diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 77f9628ba0b29..5a20bf9cf1973 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -49,8 +49,7 @@ public void testOneConsumerNoTopic() { Map topics = new HashMap<>(); topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); Map members = new HashMap<>(); - List subscribedTopics = new ArrayList<>(); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopics, new HashMap<>())); + members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), Collections.emptyList(), new HashMap<>())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment groupAssignment = assignor.assign(assignmentSpec); @@ -83,10 +82,10 @@ public void testFirstAssignmentTwoConsumersTwoTopicsSameSubscriptions() { // Members Map members = new HashMap<>(); // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + List subscribedTopicsA = Arrays.asList(topic1Uuid, topic3Uuid); members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + List subscribedTopicsB = Arrays.asList(topic1Uuid, topic3Uuid); members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); @@ -114,13 +113,13 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() // Members Map members = new HashMap<>(); // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + List subscribedTopicsA = Arrays.asList(topic1Uuid, topic2Uuid); members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); // Consumer B - List subscribedTopicsB = new ArrayList<>(Collections.singletonList(topic3Uuid)); + List subscribedTopicsB = Collections.singletonList(topic3Uuid); members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); // Consumer C - List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic2Uuid, topic3Uuid)); + List subscribedTopicsC = Arrays.asList(topic2Uuid, topic3Uuid); members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); @@ -149,13 +148,13 @@ public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { // Members Map members = new HashMap<>(); // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + List subscribedTopicsA = Arrays.asList(topic1Uuid, topic3Uuid); members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + List subscribedTopicsB = Arrays.asList(topic1Uuid, topic3Uuid); members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); // Consumer C - List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic3Uuid)); + List subscribedTopicsC = Arrays.asList(topic1Uuid, topic3Uuid); members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); From 6ced281bbebda250a58aa0646a5a88b0d895bc52 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Wed, 19 Apr 2023 15:52:33 -0700 Subject: [PATCH 23/44] addressed PR comments --- .../coordinator/group/assignor/AssignmentMemberSpec.java | 8 ++++---- .../kafka/coordinator/group/assignor/AssignmentSpec.java | 4 ++-- .../kafka/coordinator/group/assignor/RangeAssignor.java | 4 +--- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java index 6de0978c41782..4a93fccae35d5 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentMemberSpec.java @@ -115,9 +115,9 @@ public int hashCode() { @Override public String toString() { return "AssignmentMemberSpec(instanceId=" + instanceId + - ", rackId=" + rackId + - ", subscribedTopicIds=" + subscribedTopicIds + - ", assignedPartitions=" + assignedPartitions + - ')'; + ", rackId=" + rackId + + ", subscribedTopicIds=" + subscribedTopicIds + + ", assignedPartitions=" + assignedPartitions + + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java index 03fd663a40465..bf775439077ea 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignmentSpec.java @@ -78,7 +78,7 @@ public int hashCode() { @Override public String toString() { return "AssignmentSpec(members=" + members + - ", topics=" + topics + - ')'; + ", topics=" + topics + + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 8ed67361ef86d..2f45c8d3da310 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -32,7 +32,7 @@ import static java.lang.Math.min; /** - *

    The Server Side Sticky Range Assignor inherits properties of both the range assignor and the sticky assignor. + * The Server Side Sticky Range Assignor inherits properties of both the range assignor and the sticky assignor. * Properties are as follows: *

      *
    1. Each member must get at least one partition for every topic that it is subscribed to. The only exception is when @@ -46,8 +46,6 @@ * *
    2. Members should retain as much as their previous assignment as possible to reduce the number of partition movements during reassignment. (Sticky)
      3. *
    - *

    - * */ public class RangeAssignor implements PartitionAssignor { From 6aeccd6e498f818e2cb0e1499dec432142cb51a6 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Wed, 19 Apr 2023 17:09:15 -0700 Subject: [PATCH 24/44] fixed formatting stuff that was mentioned in PR comments --- .../group/assignor/RangeAssignorTest.java | 126 +++++++++++++++--- 1 file changed, 105 insertions(+), 21 deletions(-) diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 5a20bf9cf1973..b3729f68b07d9 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -49,7 +49,11 @@ public void testOneConsumerNoTopic() { Map topics = new HashMap<>(); topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); Map members = new HashMap<>(); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), Collections.emptyList(), new HashMap<>())); + members.put(consumerA, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + Collections.emptyList(), + Collections.emptyMap())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment groupAssignment = assignor.assign(assignmentSpec); @@ -64,7 +68,11 @@ public void testOneConsumerNonExistentTopic() { Map members = new HashMap<>(); List subscribedTopics = new ArrayList<>(); subscribedTopics.add(topic2Uuid); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopics, new HashMap<>())); + members.put(consumerA, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopics, + Collections.emptyMap())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment groupAssignment = assignor.assign(assignmentSpec); @@ -83,10 +91,18 @@ public void testFirstAssignmentTwoConsumersTwoTopicsSameSubscriptions() { Map members = new HashMap<>(); // Consumer A List subscribedTopicsA = Arrays.asList(topic1Uuid, topic3Uuid); - members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); + members.put(consumerA, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsA, + Collections.emptyMap())); // Consumer B List subscribedTopicsB = Arrays.asList(topic1Uuid, topic3Uuid); - members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); + members.put(consumerB, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsB, + Collections.emptyMap())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); @@ -114,13 +130,25 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() Map members = new HashMap<>(); // Consumer A List subscribedTopicsA = Arrays.asList(topic1Uuid, topic2Uuid); - members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); + members.put(consumerA, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsA, + Collections.emptyMap())); // Consumer B List subscribedTopicsB = Collections.singletonList(topic3Uuid); - members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); + members.put(consumerB, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsB, + Collections.emptyMap())); // Consumer C List subscribedTopicsC = Arrays.asList(topic2Uuid, topic3Uuid); - members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + members.put(consumerC, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsC, + Collections.emptyMap())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); @@ -149,13 +177,25 @@ public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { Map members = new HashMap<>(); // Consumer A List subscribedTopicsA = Arrays.asList(topic1Uuid, topic3Uuid); - members.put(consumerA, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, new HashMap<>())); + members.put(consumerA, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsA, + Collections.emptyMap())); // Consumer B List subscribedTopicsB = Arrays.asList(topic1Uuid, topic3Uuid); - members.put(consumerB, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, new HashMap<>())); + members.put(consumerB, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsB, + Collections.emptyMap())); // Consumer C List subscribedTopicsC = Arrays.asList(topic1Uuid, topic3Uuid); - members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + members.put(consumerC, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsC, + Collections.emptyMap())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); @@ -183,19 +223,31 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA // Add a new consumer to trigger a re-assignment // Consumer C List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + members.put(consumerC, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsC, + Collections.emptyMap())); // Consumer A List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Collections.singletonList(0))); currentAssignmentForA.put(topic2Uuid, new HashSet<>(Collections.singletonList(0))); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + members.put(consumerA, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsA, + currentAssignmentForA)); // Consumer B List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(1))); currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(1))); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + members.put(consumerB, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsB, + currentAssignmentForB)); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); @@ -227,14 +279,22 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + members.put(consumerA, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsA, + currentAssignmentForA)); // Consumer B List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + members.put(consumerB, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsB, + currentAssignmentForB)); // Simulating adding a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions Map topics = new HashMap<>(); @@ -272,17 +332,29 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + members.put(consumerA, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsA, + currentAssignmentForA)); // Consumer B List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + members.put(consumerB, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsB, + currentAssignmentForB)); // Add a new consumer to trigger a re-assignment // Consumer C List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + members.put(consumerC, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsC, + Collections.emptyMap())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); @@ -320,7 +392,11 @@ public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoC Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + members.put(consumerB, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsB, + currentAssignmentForB)); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); @@ -352,13 +428,21 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme currentAssignmentForA.put(topic2Uuid, new HashSet<>(Collections.singletonList(0))); // Change subscriptions List subscribedTopicsA = new ArrayList<>(Collections.singletonList(topic1Uuid)); - members.computeIfAbsent(consumerA, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsA, currentAssignmentForA)); + members.put(consumerA, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsA, + currentAssignmentForA)); // Consumer B Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(1))); // Change subscriptions List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid, topic3Uuid)); - members.computeIfAbsent(consumerB, k -> new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsB, currentAssignmentForB)); + members.put(consumerB, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + subscribedTopicsB, + currentAssignmentForB)); // Consumer C Map> currentAssignmentForC = new HashMap<>(); currentAssignmentForC.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); From faeda40d1a6db0b96ada5886c89a061aab6b4b97 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 20 Apr 2023 16:22:21 -0700 Subject: [PATCH 25/44] Fixed formatting and added mkAssignment function --- .../group/assignor/RangeAssignor.java | 2 +- .../group/assignor/RangeAssignorTest.java | 445 +++++++++--------- 2 files changed, 231 insertions(+), 216 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 2f45c8d3da310..5550d119b6eed 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -248,7 +248,7 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit // Consolidate the maps into MemberAssignment and then finally map each member to a MemberAssignment. Map membersWithNewAssignment = new HashMap<>(); - membersWithNewAssignmentPerTopic.forEach((memberId, assignmentPerTopic) -> membersWithNewAssignment.computeIfAbsent(memberId, k -> new MemberAssignment(assignmentPerTopic))); + membersWithNewAssignmentPerTopic.forEach((memberId, assignmentPerTopic) -> membersWithNewAssignment.put(memberId, new MemberAssignment(assignmentPerTopic))); return new GroupAssignment(membersWithNewAssignment); } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index b3729f68b07d9..46c65af2460aa 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -21,12 +21,11 @@ import org.junit.jupiter.api.Test; -import java.util.ArrayList; import java.util.Arrays; +import java.util.Collection; 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.Set; @@ -46,14 +45,15 @@ public class RangeAssignorTest { @Test public void testOneConsumerNoTopic() { - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); - Map members = new HashMap<>(); - members.put(consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - Collections.emptyList(), - Collections.emptyMap())); + Map topics = Collections.singletonMap(topic1Uuid, new AssignmentTopicMetadata(3)); + Map members = Collections.singletonMap( + consumerA, + new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + Collections.emptyList(), + Collections.emptyMap()) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment groupAssignment = assignor.assign(assignmentSpec); @@ -63,16 +63,15 @@ public void testOneConsumerNoTopic() { @Test public void testOneConsumerNonExistentTopic() { - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); - Map members = new HashMap<>(); - List subscribedTopics = new ArrayList<>(); - subscribedTopics.add(topic2Uuid); - members.put(consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopics, - Collections.emptyMap())); + Map topics = Collections.singletonMap(topic1Uuid, new AssignmentTopicMetadata(3)); + Map members = Collections.singletonMap( + consumerA, + new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + Collections.singletonList(topic2Uuid), + Collections.emptyMap()) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment groupAssignment = assignor.assign(assignmentSpec); @@ -82,187 +81,189 @@ public void testOneConsumerNonExistentTopic() { @Test public void testFirstAssignmentTwoConsumersTwoTopicsSameSubscriptions() { - // A -> T1, T3 // B -> T1, T3 // T1 -> 3 Partitions // T3 -> 2 Partitions - // Topics Map topics = new HashMap<>(); topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); - // Members + Map members = new HashMap<>(); - // Consumer A - List subscribedTopicsA = Arrays.asList(topic1Uuid, topic3Uuid); + // Initial Subscriptions are: A -> T1, T3 | B -> T1, T3 + members.put(consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsA, - Collections.emptyMap())); - // Consumer B - List subscribedTopicsB = Arrays.asList(topic1Uuid, topic3Uuid); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic3Uuid), + Collections.emptyMap()) + ); + members.put(consumerB, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsB, - Collections.emptyMap())); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic3Uuid), + Collections.emptyMap()) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1)); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(2)); // Topic 3 Partitions Assignment - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(0)); + mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(1)); assertAssignment(expectedAssignment, computedAssignment); } @Test public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() { - // A -> T1, T2 // B -> T3 // C -> T2, T3 // T1 -> 3 Partitions // T2 -> 3 Partitions // T3 -> 2 Partitions - // Topics Map topics = new HashMap<>(); topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); - // Members + Map members = new HashMap<>(); - // Consumer A - List subscribedTopicsA = Arrays.asList(topic1Uuid, topic2Uuid); + // Initial Subscriptions: A -> T1, T2 | B -> T3 | C -> T2, T3 + members.put(consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsA, - Collections.emptyMap())); - // Consumer B - List subscribedTopicsB = Collections.singletonList(topic3Uuid); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + Collections.emptyMap()) + ); + members.put(consumerB, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsB, - Collections.emptyMap())); - // Consumer C - List subscribedTopicsC = Arrays.asList(topic2Uuid, topic3Uuid); + Optional.empty(), + Optional.empty(), + Collections.singletonList(topic3Uuid), + Collections.emptyMap()) + ); + members.put(consumerC, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsC, - Collections.emptyMap())); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic2Uuid, topic3Uuid), + Collections.emptyMap()) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); + mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1, 2)); // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(0, 1)); + mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(2)); // Topic 3 Partitions Assignment - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(0)); + mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(1)); assertAssignment(expectedAssignment, computedAssignment); } @Test public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { - // Topics Map topics = new HashMap<>(); topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); - // Topic 3 has 2 partitions but three consumers subscribed to it - one of them will not get an assignment topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); - // Members + Map members = new HashMap<>(); - // Consumer A - List subscribedTopicsA = Arrays.asList(topic1Uuid, topic3Uuid); + // Initial Subscriptions: A -> T1, T3 | B -> T1, T3 | C -> T1, T3 + members.put(consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsA, - Collections.emptyMap())); - // Consumer B - List subscribedTopicsB = Arrays.asList(topic1Uuid, topic3Uuid); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic3Uuid), + Collections.emptyMap()) + ); + members.put(consumerB, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsB, - Collections.emptyMap())); - // Consumer C - List subscribedTopicsC = Arrays.asList(topic1Uuid, topic3Uuid); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic3Uuid), + Collections.emptyMap()) + ); + members.put(consumerC, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsC, - Collections.emptyMap())); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic3Uuid), + Collections.emptyMap()) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); + // Topic 3 has 2 partitions but three consumers subscribed to it - one of them will not get a partition. // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(0)); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(1)); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(2)); // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(0)); + mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(1)); assertAssignment(expectedAssignment, computedAssignment); } @Test public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerAdded() { - // Topics Map topics = new HashMap<>(); topics.put(topic1Uuid, new AssignmentTopicMetadata(2)); topics.put(topic2Uuid, new AssignmentTopicMetadata(2)); - // Members + Map members = new HashMap<>(); - // Add a new consumer to trigger a re-assignment - // Consumer C - List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); - members.put(consumerC, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsC, - Collections.emptyMap())); - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + // Initial Subscriptions: A -> T1, T2 | B -> T1, T2 | C -> T1, T2 + Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Collections.singletonList(0))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Collections.singletonList(0))); + currentAssignmentForA.put(topic1Uuid, Collections.singleton(0)); + currentAssignmentForA.put(topic2Uuid, Collections.singleton(0)); members.put(consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsA, - currentAssignmentForA)); - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForA) + ); + Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(1))); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(1))); + currentAssignmentForB.put(topic1Uuid, Collections.singleton(1)); + currentAssignmentForB.put(topic2Uuid, Collections.singleton(1)); members.put(consumerB, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsB, - currentAssignmentForB)); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForB) + ); + + // Add a new consumer to trigger a re-assignment + members.put(consumerC, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + Collections.emptyMap()) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(0)); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(1)); // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(0)); + mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(1)); // Test for stickiness - assertEquals(computedAssignment.members().get(consumerA).targetPartitions(), new HashMap<>(currentAssignmentForA), "Stickiness test failed for Consumer A"); - assertEquals(computedAssignment.members().get(consumerB).targetPartitions(), new HashMap<>(currentAssignmentForB), "Stickiness test failed for Consumer B"); + assertEquals(computedAssignment.members().get(consumerA) + .targetPartitions(), new HashMap<>(currentAssignmentForA), + "Stickiness test failed for consumer A"); + assertEquals(computedAssignment.members().get(consumerB) + .targetPartitions(), new HashMap<>(currentAssignmentForB), + "Stickiness test failed for consumer B"); // Consumer C shouldn't get any assignment, due to stickiness A, B retain their assignments assertNull(computedAssignment.members().get(consumerC)); @@ -272,49 +273,50 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA @Test public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { - // T1, T2 both have 3 partitions each when first assignment was calculated -> currentAssignmentForX + // Simulating adding a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions + Map topics = new HashMap<>(); + topics.put(topic1Uuid, new AssignmentTopicMetadata(4)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(4)); + Map members = new HashMap<>(); - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + // Initial Subscriptions: A -> T1, T2 | B -> T1, T2 + Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); members.put(consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsA, - currentAssignmentForA)); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForA) + ); - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForB.put(topic1Uuid, Collections.singleton(2)); + currentAssignmentForB.put(topic2Uuid, Collections.singleton(2)); members.put(consumerB, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsB, - currentAssignmentForB)); - - // Simulating adding a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions - Map topics = new HashMap<>(); - topics.put(topic1Uuid, new AssignmentTopicMetadata(4)); - topics.put(topic2Uuid, new AssignmentTopicMetadata(4)); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForB) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(2, 3))); + mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1)); + mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(2, 3)); // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(2, 3))); + mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(0, 1)); + mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(2, 3)); // Test for stickiness - assertEquals(computedAssignment.members().get(consumerA).targetPartitions(), new HashMap<>(currentAssignmentForA), "Stickiness test failed for Consumer A"); - // Implicitly B also retained its partitions since the other set of assigned partitions is (2,3) + assertEquals(computedAssignment.members().get(consumerA) + .targetPartitions(), new HashMap<>(currentAssignmentForA), + "Stickiness test failed for consumer A"); + // Implicitly stickiness is checked for B also since the other set of assigned partitions is (2,3) assertAssignment(expectedAssignment, computedAssignment); assertCoPartitionJoinProperty(computedAssignment); @@ -327,52 +329,55 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); Map members = new HashMap<>(); - // Consumer A - List subscribedTopicsA = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + // Initial Subscriptions: A -> T1, T2 | B -> T1, T2 | C -> T1, T2 + Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); members.put(consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsA, - currentAssignmentForA)); - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForA) + ); + Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForB.put(topic1Uuid, Collections.singleton(2)); + currentAssignmentForB.put(topic2Uuid, Collections.singleton(2)); members.put(consumerB, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsB, - currentAssignmentForB)); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForB) + ); + // Add a new consumer to trigger a re-assignment - // Consumer C - List subscribedTopicsC = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); members.put(consumerC, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsC, - Collections.emptyMap())); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + Collections.emptyMap())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(0)); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(2)); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(1)); // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(0))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(1))); + mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(0)); + mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(2)); + mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(1)); // Test for stickiness - assertEquals(computedAssignment.members().get(consumerB).targetPartitions(), new HashMap<>(currentAssignmentForB), "Stickiness test failed for Consumer B"); - assertTrue(computedAssignment.members().get(consumerA).targetPartitions().get(topic1Uuid).contains(0), "Stickiness test failed for Consumer A"); - // Co-partition join property test ensures that A retained the same partition for all topics. + assertEquals(computedAssignment.members().get(consumerB) + .targetPartitions(), new HashMap<>(currentAssignmentForB), + "Stickiness test failed for Consumer B"); + assertTrue(computedAssignment.members().get(consumerA) + .targetPartitions().get(topic1Uuid).contains(0), + "Stickiness test failed for Consumer A"); assertAssignment(expectedAssignment, computedAssignment); assertCoPartitionJoinProperty(computedAssignment); @@ -387,25 +392,24 @@ public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoC Map members = new HashMap<>(); // Consumer A was removed - // Consumer B - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid)); Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); members.put(consumerB, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsB, - currentAssignmentForB)); + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForB) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); + mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1, 2)); // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1, 2))); + mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(0, 1, 2)); assertAssignment(expectedAssignment, computedAssignment); assertCoPartitionJoinProperty(computedAssignment); @@ -419,59 +423,70 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); + Map members = new HashMap<>(); // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3 // Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2 - Map members = new HashMap<>(); - // Consumer A + Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1, 2))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Collections.singletonList(0))); - // Change subscriptions - List subscribedTopicsA = new ArrayList<>(Collections.singletonList(topic1Uuid)); + currentAssignmentForA.put(topic2Uuid, Collections.singleton(0)); members.put(consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsA, - currentAssignmentForA)); - // Consumer B + Optional.empty(), + Optional.empty(), + Collections.singletonList(topic1Uuid), + currentAssignmentForA) + ); + Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(1))); - // Change subscriptions - List subscribedTopicsB = new ArrayList<>(Arrays.asList(topic1Uuid, topic2Uuid, topic3Uuid)); + currentAssignmentForB.put(topic2Uuid, Collections.singleton(1)); members.put(consumerB, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - subscribedTopicsB, - currentAssignmentForB)); - // Consumer C + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid, topic3Uuid), + currentAssignmentForB) + ); + Map> currentAssignmentForC = new HashMap<>(); - currentAssignmentForC.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + currentAssignmentForC.put(topic2Uuid, Collections.singleton(2)); currentAssignmentForC.put(topic3Uuid, new HashSet<>(Arrays.asList(0, 1))); - // Change subscriptions - List subscribedTopicsC = new ArrayList<>(Collections.singletonList(topic2Uuid)); - members.put(consumerC, new AssignmentMemberSpec(Optional.empty(), Optional.empty(), subscribedTopicsC, new HashMap<>())); + members.put(consumerC, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + Collections.singletonList(topic2Uuid), + currentAssignmentForC) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic1Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1)); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(2)); // Topic 2 Partitions Assignment - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); - expectedAssignment.computeIfAbsent(topic2Uuid, k -> new HashSet<>()).add(new HashSet<>(Collections.singletonList(2))); + mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(0, 1)); + mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(2)); // Topic 3 Partitions Assignment - expectedAssignment.computeIfAbsent(topic3Uuid, k -> new HashSet<>()).add(new HashSet<>(Arrays.asList(0, 1))); + mkAssignment(expectedAssignment, topic3Uuid, Arrays.asList(0, 1)); // Test for stickiness - assertTrue(computedAssignment.members().get(consumerC).targetPartitions().get(topic2Uuid).contains(2), "Stickiness test failed for Consumer C"); - assertTrue(computedAssignment.members().get(consumerA).targetPartitions().get(topic1Uuid).containsAll(Arrays.asList(0, 1)), "Stickiness test failed for Consumer A"); + assertTrue(computedAssignment.members().get(consumerC) + .targetPartitions().get(topic2Uuid).contains(2), + "Stickiness test failed for Consumer C"); + assertTrue(computedAssignment.members().get(consumerA) + .targetPartitions().get(topic1Uuid).containsAll(Arrays.asList(0, 1)), + "Stickiness test failed for Consumer A"); assertAssignment(expectedAssignment, computedAssignment); } + private void mkAssignment(Map>> expectedAssignment, Uuid topicId, Collection partitions) { + expectedAssignment.computeIfAbsent(topicId, k -> new HashSet<>()).add(new HashSet<>(partitions)); + } + // We have a set of sets with the partitions that should be distributed amongst the consumers, if it exists then remove it from the set. + // The test is done like this since the order in which members are assigned partitions isn't guaranteed. We are just testing if the computed + // assignment contains the expected set of partitions, irrespective of which member got them. private void assertAssignment(Map>> expectedAssignment, GroupAssignment computedGroupAssignment) { for (MemberAssignment member : computedGroupAssignment.members().values()) { Map> computedAssignmentForMember = member.targetPartitions(); From bbbf9301347f213ed562c67c2d62d182c9630f55 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 20 Apr 2023 17:24:11 -0700 Subject: [PATCH 26/44] Made remaining formatting changes to the range assignor code --- .../group/assignor/RangeAssignor.java | 85 ++++++++++--------- 1 file changed, 45 insertions(+), 40 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 5550d119b6eed..03547ae198e69 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -36,15 +36,15 @@ * Properties are as follows: *
      *
    1. Each member must get at least one partition for every topic that it is subscribed to. The only exception is when - * the number of subscribed members is greater than the number of partitions for that topic. (Range)
      2. + * the number of subscribed members is greater than the number of partitions for that topic. (Range) *
    2. Partitions should be assigned to members in a way that facilitates the join operation when required. (Range)
      3. - * This can only be done if every member is subscribed to the same topics and the topics are co-partitioned. - * Two streams are co-partitioned if the following conditions are met: - *
        + * This can only be done if every member is subscribed to the same topics and the topics are co-partitioned. + * Two streams are co-partitioned if the following conditions are met: + *
          *
        • The keys must have the same schemas. *
        • The topics involved must have the same number of partitions. - *
        - *
      • Members should retain as much as their previous assignment as possible to reduce the number of partition movements during reassignment. (Sticky)
        • + *
      + *
    3. Members should retain as much of their previous assignment as possible to reduce the number of partition movements during reassignment. (Sticky)
      4. *
    */ public class RangeAssignor implements PartitionAssignor { @@ -87,10 +87,10 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS // Only topics that are present in both the subscribed topics list and the topic metadata should be considered for assignment. if (assignmentSpec.topics().containsKey(topicId)) { membersPerTopic - .computeIfAbsent(topicId, k -> new ArrayList<>()) - .add(memberId); + .computeIfAbsent(topicId, k -> new ArrayList<>()) + .add(memberId); } else { - log.info(memberId + " subscribed to topic " + topicId + " which doesn't exist in the topic metadata"); + log.warn(memberId + " subscribed to topic " + topicId + " which doesn't exist in the topic metadata"); } } }); @@ -105,7 +105,7 @@ private Map> getUnassignedPartitionsPerTopic(final Assignmen topicsMetadata.forEach((topicId, assignmentTopicMetadata) -> { ArrayList unassignedPartitionsForTopic = new ArrayList<>(); int numPartitions = assignmentTopicMetadata.numPartitions(); - // The partitions will be in ascending order within the list of unassigned partitions per topic. + // List of unassigned partitions per topic contains the partitions in ascending order. Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.getOrDefault(topicId, new HashSet<>()); for (int i = 0; i < numPartitions; i++) { if (!assignedStickyPartitionsForTopic.contains(i)) { @@ -119,14 +119,14 @@ private Map> getUnassignedPartitionsPerTopic(final Assignmen } /** - *

    The algorithm includes the following steps: + *

    The algorithm includes the following steps: *

      *
    1. Generate a map of membersPerTopic using the given member subscriptions.
      2. *
    2. Generate a list of members (potentiallyUnfilledMembers) that have not met the minimum required quota for assignment AND - * get a list of sticky partitions that we want to retain in the new assignment.
      3. - *
    3. Add members from the potentiallyUnfilled list to the Unfilled list if they haven't met the total required quota i.e. minimum number of partitions per member + 1 (if member is designated to receive one of the excess partitions)
      4. - *
    4. Generate a list of unassigned partitions by calculating the difference between total partitions and already assigned (sticky) partitions
      5. - *
    5. Iterate through unfilled members and assign partitions from the unassigned partitions
      6. + * get a list (assignedStickyPartitionsPerTopic) of partitions that we want to retain in the new assignment. + *
    6. Add members from the potentiallyUnfilledMembers list to the unfilledMembers list if they haven't met the total required quota i.e. minimum number of partitions per member + 1 (if member is designated to receive one of the excess partitions).
      7. + *
    7. Generate a list of unassigned partitions by calculating the difference between total partitions and already assigned (sticky) partitions.
      8. + *
    8. Iterate through the unfilled members and assign partitions from the unassigned partitions.
      9. *
    *

    */ @@ -151,41 +151,44 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit AssignmentTopicMetadata topicData = assignmentSpec.topics().get(topicId); int numPartitionsForTopic = topicData.numPartitions(); - // Idle members case : When the number of members subscribed to a topic is greater than the total number of Partitions, + // Idle members case : When the number of members subscribed to a topic is greater than the total number of partitions, // all members get assigned via the "extra partitions" logic since minRequiredQuota = 0. int minRequiredQuota = numPartitionsForTopic / membersForTopic.size(); // Each member can get only ONE extra partition per topic after receiving the minimum quota. int numMembersWithExtraPartition = numPartitionsForTopic % membersForTopic.size(); for (String memberId: membersForTopic) { - // The partitions need to be in numeric order since we want the same partition numbers from each topic - // to go to the same member in case of co-partitioned topics to facilitate joins. - Set assignedPartitionsForTopic = assignmentSpec.members().get(memberId).assignedPartitions().getOrDefault(topicId, Collections.emptySet()); + Set assignedPartitionsForTopic = assignmentSpec.members().get(memberId) + .assignedPartitions().getOrDefault(topicId, Collections.emptySet()); int currentAssignmentSize = assignedPartitionsForTopic.size(); List currentAssignmentListForTopic = new ArrayList<>(assignedPartitionsForTopic); // If there were previously assigned partitions present, we want to retain them. if (currentAssignmentSize > 0) { - // We either need to retain currentSize number of partitions when currentSize < required OR required number of partitions otherwise. + // We need to retain: + // 1) currentSize number of partitions (when currentSize < required) OR + // 2) required number of partitions (when currentSize > required). int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); + // The partitions need to be in ascending order since we want the same partition numbers from each topic + // to go to the same member to facilitate joins in case of co-partitioned topics. Collections.sort(currentAssignmentListForTopic); for (int i = 0; i < retainedPartitionsCount; i++) { assignedStickyPartitionsPerTopic - .computeIfAbsent(topicId, k -> new HashSet<>()) - .add(currentAssignmentListForTopic.get(i)); + .computeIfAbsent(topicId, k -> new HashSet<>()) + .add(currentAssignmentListForTopic.get(i)); membersWithNewAssignmentPerTopic - .computeIfAbsent(memberId, k -> new HashMap<>()) - .computeIfAbsent(topicId, k -> new HashSet<>()) - .add(currentAssignmentListForTopic.get(i)); + .computeIfAbsent(memberId, k -> new HashMap<>()) + .computeIfAbsent(topicId, k -> new HashSet<>()) + .add(currentAssignmentListForTopic.get(i)); } } // Number of partitions left to reach the minRequiredQuota. int remaining = minRequiredQuota - currentAssignmentSize; - // There are 3 cases w.r.t value of remaining - // 1) remaining < 0 this means that the member has more than the min required amount. + // There are 3 cases w.r.t the value of remaining: + // 1) remaining < 0: this means that the member has more than the min required amount. if (remaining < 0 && numMembersWithExtraPartition > 0) { // In order to remain as sticky as possible, since the order of members can be different, we want the members that already had extra // partitions to retain them if it's still required, instead of assigning the extras to the first few members directly. @@ -193,15 +196,15 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit // Since we already added the minimumRequiredQuota of partitions in the previous step (until minReq - 1), we just need to // add the extra partition that will be present at the index right after min quota was satisfied. assignedStickyPartitionsPerTopic - .computeIfAbsent(topicId, k -> new HashSet<>()) - .add(currentAssignmentListForTopic.get(minRequiredQuota)); + .computeIfAbsent(topicId, k -> new HashSet<>()) + .add(currentAssignmentListForTopic.get(minRequiredQuota)); membersWithNewAssignmentPerTopic - .computeIfAbsent(memberId, k -> new HashMap<>()) - .computeIfAbsent(topicId, k -> new HashSet<>()) - .add(currentAssignmentListForTopic.get(minRequiredQuota)); + .computeIfAbsent(memberId, k -> new HashMap<>()) + .computeIfAbsent(topicId, k -> new HashSet<>()) + .add(currentAssignmentListForTopic.get(minRequiredQuota)); } else { - // 2) If remaining = 0 it has min req partitions but there is scope for getting an extra partition later on, so it is a potentialUnfilledMember. - // 3) If remaining > 0 it doesn't even have the min required partitions, and it definitely is unfilled, so it should be added to potentialUnfilledMembers. + // 2) If remaining = 0: it has min req partitions but there is scope for getting an extra partition later on, so it is a potential unfilled member. + // 3) If remaining > 0: it doesn't even have the min required partitions, and it definitely is unfilled, so it should be added to potentiallyUnfilledMembers. RemainingAssignmentsForMember newPair = new RemainingAssignmentsForMember<>(memberId, remaining); potentiallyUnfilledMembers.add(newPair); } @@ -219,8 +222,8 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit if (remaining > 0) { RemainingAssignmentsForMember newPair = new RemainingAssignmentsForMember<>(memberId, remaining); unfilledMembersPerTopic - .computeIfAbsent(topicId, k -> new ArrayList<>()) - .add(newPair); + .computeIfAbsent(topicId, k -> new ArrayList<>()) + .add(newPair); } } }); @@ -236,12 +239,14 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit for (RemainingAssignmentsForMember remainingAssignmentsForMember : unfilledMembersForTopic) { String memberId = remainingAssignmentsForMember.memberId(); int remaining = remainingAssignmentsForMember.remaining(); - List partitionsToAssign = unassignedPartitionsPerTopic.get(topicId).subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); + List partitionsToAssign = unassignedPartitionsPerTopic.get(topicId) + .subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); + unassignedPartitionsListStartPointer += remaining; membersWithNewAssignmentPerTopic - .computeIfAbsent(memberId, k -> new HashMap<>()) - .computeIfAbsent(topicId, k -> new HashSet<>()) - .addAll(partitionsToAssign); + .computeIfAbsent(memberId, k -> new HashMap<>()) + .computeIfAbsent(topicId, k -> new HashSet<>()) + .addAll(partitionsToAssign); } }); From 0d6d127bd505d697ff3761ca7e283b271c8ef903 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 25 Apr 2023 16:52:34 -0700 Subject: [PATCH 27/44] moved getUnassignedPartitions and assign unfilled members into the same topicId loop, edited comments and java doc --- .../group/assignor/RangeAssignor.java | 123 +++++++++--------- 1 file changed, 60 insertions(+), 63 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 03547ae198e69..9b41f20273e2d 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -33,22 +33,21 @@ /** * The Server Side Sticky Range Assignor inherits properties of both the range assignor and the sticky assignor. - * Properties are as follows: + * The properties are as follows: *
      - *
    1. Each member must get at least one partition for every topic that it is subscribed to. The only exception is when + *
    2. Each member must get at least one partition from every topic that it is subscribed to. The only exception is when * the number of subscribed members is greater than the number of partitions for that topic. (Range)
      3. *
    3. Partitions should be assigned to members in a way that facilitates the join operation when required. (Range)
      4. * This can only be done if every member is subscribed to the same topics and the topics are co-partitioned. * Two streams are co-partitioned if the following conditions are met: *
        - *
      • The keys must have the same schemas. - *
      • The topics involved must have the same number of partitions. + *
      • The keys must have the same schemas.
        • + *
      • The topics involved must have the same number of partitions.
        • *
      *
    4. Members should retain as much of their previous assignment as possible to reduce the number of partition movements during reassignment. (Sticky)
      5. *
    */ public class RangeAssignor implements PartitionAssignor { - private static final Logger log = LoggerFactory.getLogger(RangeAssignor.class); public static final String RANGE_ASSIGNOR_NAME = "range"; @@ -58,23 +57,31 @@ public String name() { return RANGE_ASSIGNOR_NAME; } - static class RemainingAssignmentsForMember { + // Used in the potentiallyUnfilledMembers map and the UnfilledMembers map. + private static class MemberWithRemainingAssignments { private final String memberId; + /** + * Number of partitions required to meet the assignment quota + */ private final Integer remaining; - public RemainingAssignmentsForMember(String memberId, Integer remaining) { + public MemberWithRemainingAssignments(String memberId, Integer remaining) { this.memberId = memberId; this.remaining = remaining; } + /** + * @return memberId + */ public String memberId() { return memberId; } - + /** + * @return Remaining number of partitions + */ public Integer remaining() { return remaining; } - } private Map> membersPerTopic(final AssignmentSpec assignmentSpec) { @@ -98,35 +105,17 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS return membersPerTopic; } - private Map> getUnassignedPartitionsPerTopic(final AssignmentSpec assignmentSpec, Map> assignedStickyPartitionsPerTopic) { - Map> unassignedPartitionsPerTopic = new HashMap<>(); - Map topicsMetadata = assignmentSpec.topics(); - - topicsMetadata.forEach((topicId, assignmentTopicMetadata) -> { - ArrayList unassignedPartitionsForTopic = new ArrayList<>(); - int numPartitions = assignmentTopicMetadata.numPartitions(); - // List of unassigned partitions per topic contains the partitions in ascending order. - Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.getOrDefault(topicId, new HashSet<>()); - for (int i = 0; i < numPartitions; i++) { - if (!assignedStickyPartitionsForTopic.contains(i)) { - unassignedPartitionsForTopic.add(i); - } - } - unassignedPartitionsPerTopic.put(topicId, unassignedPartitionsForTopic); - }); - - return unassignedPartitionsPerTopic; - } - /** *

    The algorithm includes the following steps: *

      *
    1. Generate a map of membersPerTopic using the given member subscriptions.
      2. - *
    2. Generate a list of members (potentiallyUnfilledMembers) that have not met the minimum required quota for assignment AND - * get a list (assignedStickyPartitionsPerTopic) of partitions that we want to retain in the new assignment.
      3. - *
    3. Add members from the potentiallyUnfilledMembers list to the unfilledMembers list if they haven't met the total required quota i.e. minimum number of partitions per member + 1 (if member is designated to receive one of the excess partitions).
      4. - *
    4. Generate a list of unassigned partitions by calculating the difference between total partitions and already assigned (sticky) partitions.
      5. - *
    5. Iterate through the unfilled members and assign partitions from the unassigned partitions.
      6. + *
    6. Generate a list of members (potentiallyUnfilledMembers) that have not met the minimum required quota of partitions for the assignment AND + * get a list (assignedStickyPartitionsPerTopic) of partitions that will be retained in the new assignment.
      7. + *
    7. Add members from the potentiallyUnfilledMembers list to the unfilledMembersPerTopic map if they haven't met the total required quota + * i.e. minRequiredQuota + 1, if the member is designated to receive one of the excess partitions OR minRequiredQuota otherwise.
      8. + *
    8. Generate a list of unassigned partitions by calculating the difference between the total partitions for the topic and the assigned (sticky) partitions.
      9. + *
    9. Check if unfilled members exist for the current topicId and assign partitions to them in ranges from the unassignedPartitionsPerTopic map + * based on the remaining partitions value stored.
      10. *
    *

    */ @@ -136,8 +125,9 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit // Step 1 Map> membersPerTopic = membersPerTopic(assignmentSpec); // Step 2 - Map>> unfilledMembersPerTopic = new HashMap<>(); + Map> unfilledMembersPerTopic = new HashMap<>(); Map> assignedStickyPartitionsPerTopic = new HashMap<>(); + Map> unassignedPartitionsPerTopic = new HashMap<>(); membersPerTopic.forEach((topicId, membersForTopic) -> { // For each topic we have a temporary list of members stored in potentiallyUnfilledMembers. @@ -146,7 +136,7 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit // BUT they could be assigned an extra partition later on. Extra partitions exist if total partitions % number of members is greater than 0. // In this case we add the member to the unfilled members map iff an extra partition needs to be assigned to it. // 2) Members that don't have the minimum required partitions, so irrespective of whether they get an extra partition or not they get added to the unfilled map later. - List> potentiallyUnfilledMembers = new ArrayList<>(); + List potentiallyUnfilledMembers = new ArrayList<>(); AssignmentTopicMetadata topicData = assignmentSpec.topics().get(topicId); int numPartitionsForTopic = topicData.numPartitions(); @@ -184,7 +174,7 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit } } - // Number of partitions left to reach the minRequiredQuota. + // Number of partitions required to meet the minRequiredQuota. int remaining = minRequiredQuota - currentAssignmentSize; // There are 3 cases w.r.t the value of remaining: @@ -203,16 +193,16 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit .computeIfAbsent(topicId, k -> new HashSet<>()) .add(currentAssignmentListForTopic.get(minRequiredQuota)); } else { - // 2) If remaining = 0: it has min req partitions but there is scope for getting an extra partition later on, so it is a potential unfilled member. - // 3) If remaining > 0: it doesn't even have the min required partitions, and it definitely is unfilled, so it should be added to potentiallyUnfilledMembers. - RemainingAssignmentsForMember newPair = new RemainingAssignmentsForMember<>(memberId, remaining); + // 2) If remaining = 0: member has the minimum required partitions, but it may get an extra partition, so it is a potentially unfilled member. + // 3) If remaining > 0: member doesn't have the minimum required partitions, so it should be added to potentiallyUnfilledMembers. + MemberWithRemainingAssignments newPair = new MemberWithRemainingAssignments(memberId, remaining); potentiallyUnfilledMembers.add(newPair); } } // Step 3 // If remaining > 0 after increasing the required quota due to the extra partition, add potentially unfilled member to the unfilled members list. - for (RemainingAssignmentsForMember pair : potentiallyUnfilledMembers) { + for (MemberWithRemainingAssignments pair : potentiallyUnfilledMembers) { String memberId = pair.memberId(); Integer remaining = pair.remaining(); if (numMembersWithExtraPartition > 0) { @@ -220,41 +210,48 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit numMembersWithExtraPartition--; } if (remaining > 0) { - RemainingAssignmentsForMember newPair = new RemainingAssignmentsForMember<>(memberId, remaining); + MemberWithRemainingAssignments newPair = new MemberWithRemainingAssignments(memberId, remaining); unfilledMembersPerTopic .computeIfAbsent(topicId, k -> new ArrayList<>()) .add(newPair); } } - }); - - // Step 4 - // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get unassigned partitions. - Map> unassignedPartitionsPerTopic = getUnassignedPartitionsPerTopic(assignmentSpec, assignedStickyPartitionsPerTopic); - // Step 5 - // Iterate through the unfilled members list and assign the remaining number of partitions from the unassignedPartitions list. - unfilledMembersPerTopic.forEach((topicId, unfilledMembersForTopic) -> { - int unassignedPartitionsListStartPointer = 0; - for (RemainingAssignmentsForMember remainingAssignmentsForMember : unfilledMembersForTopic) { - String memberId = remainingAssignmentsForMember.memberId(); - int remaining = remainingAssignmentsForMember.remaining(); - List partitionsToAssign = unassignedPartitionsPerTopic.get(topicId) - .subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); + // Step 4 + // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get unassigned partitions. + ArrayList unassignedPartitionsForTopic = new ArrayList<>(); + int numPartitions = assignmentSpec.topics().get(topicId).numPartitions(); + // List of unassigned partitions per topic contains the partitions in ascending order. + Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.getOrDefault(topicId, new HashSet<>()); + for (int i = 0; i < numPartitions; i++) { + if (!assignedStickyPartitionsForTopic.contains(i)) { + unassignedPartitionsForTopic.add(i); + } + } + unassignedPartitionsPerTopic.put(topicId, unassignedPartitionsForTopic); - unassignedPartitionsListStartPointer += remaining; - membersWithNewAssignmentPerTopic - .computeIfAbsent(memberId, k -> new HashMap<>()) - .computeIfAbsent(topicId, k -> new HashSet<>()) - .addAll(partitionsToAssign); + // Step 5 + // Assign the remaining number of partitions from the unassignedPartitions list if unfilled members exist for this topic. + if (unfilledMembersPerTopic.containsKey(topicId)) { + int unassignedPartitionsListStartPointer = 0; + for (MemberWithRemainingAssignments memberWithRemainingAssignments : unfilledMembersPerTopic.get(topicId)) { + String memberId = memberWithRemainingAssignments.memberId(); + int remaining = memberWithRemainingAssignments.remaining(); + List partitionsToAssign = unassignedPartitionsPerTopic.get(topicId) + .subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); + + unassignedPartitionsListStartPointer += remaining; + membersWithNewAssignmentPerTopic + .computeIfAbsent(memberId, k -> new HashMap<>()) + .computeIfAbsent(topicId, k -> new HashSet<>()) + .addAll(partitionsToAssign); + } } }); // Consolidate the maps into MemberAssignment and then finally map each member to a MemberAssignment. Map membersWithNewAssignment = new HashMap<>(); - membersWithNewAssignmentPerTopic.forEach((memberId, assignmentPerTopic) -> membersWithNewAssignment.put(memberId, new MemberAssignment(assignmentPerTopic))); - return new GroupAssignment(membersWithNewAssignment); } } From 852be6b9f8134cf82befcda3ea759f7613e7d06c Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 25 Apr 2023 16:58:56 -0700 Subject: [PATCH 28/44] suppression removed --- checkstyle/suppressions.xml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index ec8f9b2e9d892..28cefdd30c2b8 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -318,12 +318,6 @@ - - - - From 0a3dc7c7703ad056fd792dd63bd19d2ba7dbf7b4 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 25 Apr 2023 16:59:51 -0700 Subject: [PATCH 29/44] suppression removed --- checkstyle/suppressions.xml | 1 - 1 file changed, 1 deletion(-) diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index d0a10e6b388ce..af271ee77bd3b 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -318,7 +318,6 @@ - From 0823f2a5ab2abbba34909dee8640ef06ea0f156f Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 25 Apr 2023 17:14:21 -0700 Subject: [PATCH 30/44] added another test to check if stickiness is retained with the extra partition and when there are multiple triggers for reassignment --- .../group/assignor/RangeAssignorTest.java | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 46c65af2460aa..ebcc42bed1140 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -382,6 +382,66 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon assertAssignment(expectedAssignment, computedAssignment); assertCoPartitionJoinProperty(computedAssignment); } + @Test + public void testReassignmentWhenOneConsumerAddedAndOnePartitionAfterInitialAssignmentWithTwoConsumersTwoTopics() { + Map topics = new HashMap<>(); + // Add a new partition to topic 1, initially T1 -> 3 partitions + topics.put(topic1Uuid, new AssignmentTopicMetadata(4)); + topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); + + Map members = new HashMap<>(); + // Initial Subscriptions: A -> T1, T2 | B -> T1, T2 | C -> T1, T2 + + Map> currentAssignmentForA = new HashMap<>(); + currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); + currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); + members.put(consumerA, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForA) + ); + + Map> currentAssignmentForB = new HashMap<>(); + currentAssignmentForB.put(topic1Uuid, Collections.singleton(2)); + currentAssignmentForB.put(topic2Uuid, Collections.singleton(2)); + members.put(consumerB, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForB) + ); + + // Add a new consumer to trigger a re-assignment + members.put(consumerC, new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + Collections.singletonList(topic1Uuid), + Collections.emptyMap())); + + AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); + GroupAssignment computedAssignment = assignor.assign(assignmentSpec); + + Map>> expectedAssignment = new HashMap<>(); + // Topic 1 Partitions Assignment + mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1)); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(2)); + mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(3)); + // Topic 2 Partitions Assignment + // Since the new consumer isn't subscribed to topic 2 the assignment shouldn't change + mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(0, 1)); + mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(2)); + System.out.println("assignment is " + computedAssignment); + // Test for stickiness + assertEquals(computedAssignment.members().get(consumerA) + .targetPartitions(), new HashMap<>(currentAssignmentForA), + "Stickiness test failed for Consumer A"); + assertEquals(computedAssignment.members().get(consumerB) + .targetPartitions(), new HashMap<>(currentAssignmentForB), + "Stickiness test failed for Consumer B"); + + assertAssignment(expectedAssignment, computedAssignment); + } @Test public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoConsumersTwoTopics() { From 91f941795bec6e54da21e465aa14bab6dc1a76dd Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 25 Apr 2023 17:18:28 -0700 Subject: [PATCH 31/44] nit --- .../apache/kafka/coordinator/group/assignor/RangeAssignor.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 9b41f20273e2d..67ae061ceda88 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -97,7 +97,7 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS .computeIfAbsent(topicId, k -> new ArrayList<>()) .add(memberId); } else { - log.warn(memberId + " subscribed to topic " + topicId + " which doesn't exist in the topic metadata"); + log.warn("Member " + memberId + " subscribed to topic " + topicId + " which doesn't exist in the topic metadata"); } } }); From 078cb9f8fc5273326d29f1ffc1abdf6373db3aec Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 27 Apr 2023 08:12:50 -0700 Subject: [PATCH 32/44] indentation fixes --- .../group/assignor/RangeAssignor.java | 6 +-- .../group/assignor/RangeAssignorTest.java | 40 +++++++++---------- 2 files changed, 23 insertions(+), 23 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 67ae061ceda88..58dcd625254db 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -242,9 +242,9 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit unassignedPartitionsListStartPointer += remaining; membersWithNewAssignmentPerTopic - .computeIfAbsent(memberId, k -> new HashMap<>()) - .computeIfAbsent(topicId, k -> new HashSet<>()) - .addAll(partitionsToAssign); + .computeIfAbsent(memberId, k -> new HashMap<>()) + .computeIfAbsent(topicId, k -> new HashSet<>()) + .addAll(partitionsToAssign); } } }); diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index ebcc42bed1140..3fefa7b4b13bf 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -49,10 +49,10 @@ public void testOneConsumerNoTopic() { Map members = Collections.singletonMap( consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - Collections.emptyList(), - Collections.emptyMap()) + Optional.empty(), + Optional.empty(), + Collections.emptyList(), + Collections.emptyMap()) ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); @@ -240,10 +240,10 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA // Add a new consumer to trigger a re-assignment members.put(consumerC, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - Arrays.asList(topic1Uuid, topic2Uuid), - Collections.emptyMap()) + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + Collections.emptyMap()) ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); @@ -396,28 +396,28 @@ public void testReassignmentWhenOneConsumerAddedAndOnePartitionAfterInitialAssig currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); members.put(consumerA, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForA) + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForA) ); Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic1Uuid, Collections.singleton(2)); currentAssignmentForB.put(topic2Uuid, Collections.singleton(2)); members.put(consumerB, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForB) + Optional.empty(), + Optional.empty(), + Arrays.asList(topic1Uuid, topic2Uuid), + currentAssignmentForB) ); // Add a new consumer to trigger a re-assignment members.put(consumerC, new AssignmentMemberSpec( - Optional.empty(), - Optional.empty(), - Collections.singletonList(topic1Uuid), - Collections.emptyMap())); + Optional.empty(), + Optional.empty(), + Collections.singletonList(topic1Uuid), + Collections.emptyMap())); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); From 9553b103d83897a67d670156a5cf01d391540a3c Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 27 Apr 2023 16:11:59 -0700 Subject: [PATCH 33/44] Range assignor main file changes --- .../group/assignor/RangeAssignor.java | 162 ++++++++---------- .../group/assignor/RangeAssignorTest.java | 2 - 2 files changed, 71 insertions(+), 93 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 58dcd625254db..a07aa30d5703f 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -32,18 +32,19 @@ import static java.lang.Math.min; /** - * The Server Side Sticky Range Assignor inherits properties of both the range assignor and the sticky assignor. + * This Range Assignor inherits properties of both the range assignor and the sticky assignor. * The properties are as follows: *
      *
    1. Each member must get at least one partition from every topic that it is subscribed to. The only exception is when * the number of subscribed members is greater than the number of partitions for that topic. (Range)
      2. - *
    2. Partitions should be assigned to members in a way that facilitates the join operation when required. (Range)
      3. + *
    3. Partitions should be assigned to members in a way that facilitates the join operation when required. (Range) * This can only be done if every member is subscribed to the same topics and the topics are co-partitioned. * Two streams are co-partitioned if the following conditions are met: *
        *
      • The keys must have the same schemas.
        • *
      • The topics involved must have the same number of partitions.
        • *
      + *
      4. *
    4. Members should retain as much of their previous assignment as possible to reduce the number of partition movements during reassignment. (Sticky)
      5. *
    */ @@ -57,11 +58,16 @@ public String name() { return RANGE_ASSIGNOR_NAME; } - // Used in the potentiallyUnfilledMembers map and the UnfilledMembers map. + /** + * Pair of memberId and remaining partitions to meet the quota. + */ private static class MemberWithRemainingAssignments { + /** + * Member Id. + */ private final String memberId; /** - * Number of partitions required to meet the assignment quota + * Number of partitions required to meet the assignment quota. */ private final Integer remaining; @@ -69,21 +75,11 @@ public MemberWithRemainingAssignments(String memberId, Integer remaining) { this.memberId = memberId; this.remaining = remaining; } - - /** - * @return memberId - */ - public String memberId() { - return memberId; - } - /** - * @return Remaining number of partitions - */ - public Integer remaining() { - return remaining; - } } + /** + * @return Map of topicIds to a list of members subscribed to them. + */ private Map> membersPerTopic(final AssignmentSpec assignmentSpec) { Map> membersPerTopic = new HashMap<>(); Map membersData = assignmentSpec.members(); @@ -105,154 +101,138 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS return membersPerTopic; } -/** - *

    The algorithm includes the following steps: - *

      - *
    1. Generate a map of membersPerTopic using the given member subscriptions.
      2. - *
    2. Generate a list of members (potentiallyUnfilledMembers) that have not met the minimum required quota of partitions for the assignment AND - * get a list (assignedStickyPartitionsPerTopic) of partitions that will be retained in the new assignment.
      3. - *
    3. Add members from the potentiallyUnfilledMembers list to the unfilledMembersPerTopic map if they haven't met the total required quota - * i.e. minRequiredQuota + 1, if the member is designated to receive one of the excess partitions OR minRequiredQuota otherwise.
      4. - *
    4. Generate a list of unassigned partitions by calculating the difference between the total partitions for the topic and the assigned (sticky) partitions.
      5. - *
    5. Check if unfilled members exist for the current topicId and assign partitions to them in ranges from the unassignedPartitionsPerTopic map - * based on the remaining partitions value stored.
      6. - *
    - *

    - */ + /** + *

    The algorithm includes the following steps: + *

      + *
    1. Generate a map of membersPerTopic using the given member subscriptions.
      2. + *
    2. Generate a list of members (potentiallyUnfilledMembers) that have not met the minimum required quota of partitions for the assignment AND + * get a list (assignedStickyPartitionsPerTopic) of partitions that will be retained in the new assignment.
      3. + *
    3. Add members from the potentiallyUnfilledMembers list to the unfilledMembersPerTopic map if they haven't met the total required quota + * i.e. minRequiredQuota + 1, if the member is designated to receive one of the excess partitions OR minRequiredQuota otherwise.
      4. + *
    4. Generate a list of unassigned partitions by calculating the difference between the total partitions for the topic and the assigned (sticky) partitions.
      5. + *
    5. Check if unfilled members exist for the current topicId and assign partitions to them in ranges from the unassignedPartitionsPerTopic map + * based on the remaining partitions value stored.
      6. + *
    + *

    + */ @Override public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws PartitionAssignorException { - Map>> membersWithNewAssignmentPerTopic = new HashMap<>(); + Map newAssignment = new HashMap<>(); + // Step 1 Map> membersPerTopic = membersPerTopic(assignmentSpec); - // Step 2 - Map> unfilledMembersPerTopic = new HashMap<>(); - Map> assignedStickyPartitionsPerTopic = new HashMap<>(); - Map> unassignedPartitionsPerTopic = new HashMap<>(); membersPerTopic.forEach((topicId, membersForTopic) -> { - // For each topic we have a temporary list of members stored in potentiallyUnfilledMembers. - // The list is populated with members that satisfy one of the two conditions: - // 1) Members that already have the minimum required number of partitions i.e. total partitions divided by total subscribed members - // BUT they could be assigned an extra partition later on. Extra partitions exist if total partitions % number of members is greater than 0. - // In this case we add the member to the unfilled members map iff an extra partition needs to be assigned to it. - // 2) Members that don't have the minimum required partitions, so irrespective of whether they get an extra partition or not they get added to the unfilled map later. - List potentiallyUnfilledMembers = new ArrayList<>(); AssignmentTopicMetadata topicData = assignmentSpec.topics().get(topicId); int numPartitionsForTopic = topicData.numPartitions(); - - // Idle members case : When the number of members subscribed to a topic is greater than the total number of partitions, - // all members get assigned via the "extra partitions" logic since minRequiredQuota = 0. int minRequiredQuota = numPartitionsForTopic / membersForTopic.size(); // Each member can get only ONE extra partition per topic after receiving the minimum quota. int numMembersWithExtraPartition = numPartitionsForTopic % membersForTopic.size(); - for (String memberId: membersForTopic) { + // Idle members case : When the number of members subscribed to a topic is greater than the total number of partitions, + // all members get assigned via the "extra partitions" logic since minRequiredQuota = 0. + + // Step 2 + Set assignedStickyPartitionsForTopic = new HashSet<>(); + List potentiallyUnfilledMembers = new ArrayList<>(); + + for (String memberId : membersForTopic) { Set assignedPartitionsForTopic = assignmentSpec.members().get(memberId) .assignedPartitions().getOrDefault(topicId, Collections.emptySet()); int currentAssignmentSize = assignedPartitionsForTopic.size(); List currentAssignmentListForTopic = new ArrayList<>(assignedPartitionsForTopic); - // If there were previously assigned partitions present, we want to retain them. + // If there were partitions from this topic that were previously assigned to this member present, retain as much as possible. + // We need to retain: + // 1) currentSize number of partitions (when currentSize < required) OR + // 2) required number of partitions (when currentSize > required). if (currentAssignmentSize > 0) { - // We need to retain: - // 1) currentSize number of partitions (when currentSize < required) OR - // 2) required number of partitions (when currentSize > required). int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); // The partitions need to be in ascending order since we want the same partition numbers from each topic // to go to the same member to facilitate joins in case of co-partitioned topics. Collections.sort(currentAssignmentListForTopic); for (int i = 0; i < retainedPartitionsCount; i++) { - assignedStickyPartitionsPerTopic - .computeIfAbsent(topicId, k -> new HashSet<>()) + assignedStickyPartitionsForTopic .add(currentAssignmentListForTopic.get(i)); - membersWithNewAssignmentPerTopic - .computeIfAbsent(memberId, k -> new HashMap<>()) + newAssignment.computeIfAbsent(memberId, k -> new MemberAssignment(new HashMap<>())) + .targetPartitions() .computeIfAbsent(topicId, k -> new HashSet<>()) .add(currentAssignmentListForTopic.get(i)); } } // Number of partitions required to meet the minRequiredQuota. - int remaining = minRequiredQuota - currentAssignmentSize; - // There are 3 cases w.r.t the value of remaining: // 1) remaining < 0: this means that the member has more than the min required amount. + // 2) If remaining = 0: member has the minimum required partitions, but it may get an extra partition, so it is a potentially unfilled member. + // 3) If remaining > 0: member doesn't have the minimum required partitions, so it should be added to potentiallyUnfilledMembers. + int remaining = minRequiredQuota - currentAssignmentSize; + + // In order to remain as sticky as possible, since the order of members can be different, we want the members that have more than the required number of partitions, + // to get the extra partition instead of assigning the extras to the first few members directly. if (remaining < 0 && numMembersWithExtraPartition > 0) { - // In order to remain as sticky as possible, since the order of members can be different, we want the members that already had extra - // partitions to retain them if it's still required, instead of assigning the extras to the first few members directly. numMembersWithExtraPartition--; // Since we already added the minimumRequiredQuota of partitions in the previous step (until minReq - 1), we just need to // add the extra partition that will be present at the index right after min quota was satisfied. - assignedStickyPartitionsPerTopic - .computeIfAbsent(topicId, k -> new HashSet<>()) + assignedStickyPartitionsForTopic .add(currentAssignmentListForTopic.get(minRequiredQuota)); - membersWithNewAssignmentPerTopic - .computeIfAbsent(memberId, k -> new HashMap<>()) + newAssignment.computeIfAbsent(memberId, k -> new MemberAssignment(new HashMap<>())) + .targetPartitions() .computeIfAbsent(topicId, k -> new HashSet<>()) .add(currentAssignmentListForTopic.get(minRequiredQuota)); } else { - // 2) If remaining = 0: member has the minimum required partitions, but it may get an extra partition, so it is a potentially unfilled member. - // 3) If remaining > 0: member doesn't have the minimum required partitions, so it should be added to potentiallyUnfilledMembers. MemberWithRemainingAssignments newPair = new MemberWithRemainingAssignments(memberId, remaining); potentiallyUnfilledMembers.add(newPair); } } + List unfilledMembersForTopic = new ArrayList<>(); + // Step 3 // If remaining > 0 after increasing the required quota due to the extra partition, add potentially unfilled member to the unfilled members list. for (MemberWithRemainingAssignments pair : potentiallyUnfilledMembers) { - String memberId = pair.memberId(); - Integer remaining = pair.remaining(); + String memberId = pair.memberId; + Integer remaining = pair.remaining; if (numMembersWithExtraPartition > 0) { remaining++; numMembersWithExtraPartition--; } if (remaining > 0) { MemberWithRemainingAssignments newPair = new MemberWithRemainingAssignments(memberId, remaining); - unfilledMembersPerTopic - .computeIfAbsent(topicId, k -> new ArrayList<>()) + unfilledMembersForTopic .add(newPair); } } // Step 4 // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get unassigned partitions. - ArrayList unassignedPartitionsForTopic = new ArrayList<>(); - int numPartitions = assignmentSpec.topics().get(topicId).numPartitions(); // List of unassigned partitions per topic contains the partitions in ascending order. - Set assignedStickyPartitionsForTopic = assignedStickyPartitionsPerTopic.getOrDefault(topicId, new HashSet<>()); - for (int i = 0; i < numPartitions; i++) { + List unassignedPartitionsForTopic = new ArrayList<>(); + for (int i = 0; i < numPartitionsForTopic; i++) { if (!assignedStickyPartitionsForTopic.contains(i)) { unassignedPartitionsForTopic.add(i); } } - unassignedPartitionsPerTopic.put(topicId, unassignedPartitionsForTopic); // Step 5 - // Assign the remaining number of partitions from the unassignedPartitions list if unfilled members exist for this topic. - if (unfilledMembersPerTopic.containsKey(topicId)) { - int unassignedPartitionsListStartPointer = 0; - for (MemberWithRemainingAssignments memberWithRemainingAssignments : unfilledMembersPerTopic.get(topicId)) { - String memberId = memberWithRemainingAssignments.memberId(); - int remaining = memberWithRemainingAssignments.remaining(); - List partitionsToAssign = unassignedPartitionsPerTopic.get(topicId) - .subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); - - unassignedPartitionsListStartPointer += remaining; - membersWithNewAssignmentPerTopic - .computeIfAbsent(memberId, k -> new HashMap<>()) + // Assign the remaining number of partitions from the unassigned partitions list. + int unassignedPartitionsListStartPointer = 0; + for (MemberWithRemainingAssignments memberWithRemainingAssignments : unfilledMembersForTopic) { + String memberId = memberWithRemainingAssignments.memberId; + int remaining = memberWithRemainingAssignments.remaining; + List partitionsToAssign = unassignedPartitionsForTopic + .subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); + unassignedPartitionsListStartPointer += remaining; + newAssignment.computeIfAbsent(memberId, k -> new MemberAssignment(new HashMap<>())) + .targetPartitions() .computeIfAbsent(topicId, k -> new HashSet<>()) .addAll(partitionsToAssign); - } } }); - // Consolidate the maps into MemberAssignment and then finally map each member to a MemberAssignment. - Map membersWithNewAssignment = new HashMap<>(); - membersWithNewAssignmentPerTopic.forEach((memberId, assignmentPerTopic) -> membersWithNewAssignment.put(memberId, new MemberAssignment(assignmentPerTopic))); - return new GroupAssignment(membersWithNewAssignment); + return new GroupAssignment(newAssignment); } } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 3fefa7b4b13bf..65ea649ad33f4 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -14,7 +14,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package org.apache.kafka.coordinator.group.assignor; import org.apache.kafka.common.Uuid; @@ -477,7 +476,6 @@ public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoC @Test public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignmentWithThreeConsumersTwoTopics() { - Map topics = new HashMap<>(); topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); From bfcb8975a89e2bb38ba9dbf529eab9b19006ecdf Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Thu, 27 Apr 2023 16:42:09 -0700 Subject: [PATCH 34/44] Range assignor test file changes --- .../group/assignor/RangeAssignorTest.java | 33 +++++++++---------- 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 65ea649ad33f4..86a13fced762f 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -28,6 +28,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.TreeMap; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; @@ -51,17 +52,18 @@ public void testOneConsumerNoTopic() { Optional.empty(), Optional.empty(), Collections.emptyList(), - Collections.emptyMap()) - ); + Collections.emptyMap() + ) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment groupAssignment = assignor.assign(assignmentSpec); - assertTrue(groupAssignment.members().isEmpty()); + assertEquals(Collections.emptyMap(), groupAssignment.members()); } @Test - public void testOneConsumerNonExistentTopic() { + public void testOneConsumerSubscribedToNonExistentTopic() { Map topics = Collections.singletonMap(topic1Uuid, new AssignmentTopicMetadata(3)); Map members = Collections.singletonMap( consumerA, @@ -69,8 +71,9 @@ public void testOneConsumerNonExistentTopic() { Optional.empty(), Optional.empty(), Collections.singletonList(topic2Uuid), - Collections.emptyMap()) - ); + Collections.emptyMap() + ) + ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment groupAssignment = assignor.assign(assignmentSpec); @@ -84,8 +87,7 @@ public void testFirstAssignmentTwoConsumersTwoTopicsSameSubscriptions() { topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); - Map members = new HashMap<>(); - // Initial Subscriptions are: A -> T1, T3 | B -> T1, T3 + Map members = new TreeMap<>(); members.put(consumerA, new AssignmentMemberSpec( Optional.empty(), @@ -122,8 +124,7 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); - Map members = new HashMap<>(); - // Initial Subscriptions: A -> T1, T2 | B -> T3 | C -> T2, T3 + Map members = new TreeMap<>(); members.put(consumerA, new AssignmentMemberSpec( Optional.empty(), @@ -168,8 +169,7 @@ public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); - Map members = new HashMap<>(); - // Initial Subscriptions: A -> T1, T3 | B -> T1, T3 | C -> T1, T3 + Map members = new TreeMap<>(); members.put(consumerA, new AssignmentMemberSpec( Optional.empty(), @@ -214,8 +214,7 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA topics.put(topic1Uuid, new AssignmentTopicMetadata(2)); topics.put(topic2Uuid, new AssignmentTopicMetadata(2)); - Map members = new HashMap<>(); - // Initial Subscriptions: A -> T1, T2 | B -> T1, T2 | C -> T1, T2 + Map members = new TreeMap<>(); Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, Collections.singleton(0)); @@ -277,8 +276,7 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { topics.put(topic1Uuid, new AssignmentTopicMetadata(4)); topics.put(topic2Uuid, new AssignmentTopicMetadata(4)); - Map members = new HashMap<>(); - // Initial Subscriptions: A -> T1, T2 | B -> T1, T2 + Map members = new TreeMap<>(); Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); @@ -328,7 +326,6 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); Map members = new HashMap<>(); - // Initial Subscriptions: A -> T1, T2 | B -> T1, T2 | C -> T1, T2 Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); @@ -389,7 +386,6 @@ public void testReassignmentWhenOneConsumerAddedAndOnePartitionAfterInitialAssig topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); Map members = new HashMap<>(); - // Initial Subscriptions: A -> T1, T2 | B -> T1, T2 | C -> T1, T2 Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); @@ -482,6 +478,7 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); Map members = new HashMap<>(); + // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3 // Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2 From f3e72e5a69b4450d9d9f50fc8a8b046666880fdd Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Mon, 8 May 2023 17:31:17 -0700 Subject: [PATCH 35/44] Range assignor test file changes --- .../{consumer => }/AssignmentTestUtil.java | 2 +- .../group/assignor/RangeAssignorTest.java | 260 +++++++----------- .../group/consumer/AssignmentTest.java | 4 +- .../consumer/ConsumerGroupMemberTest.java | 4 +- 4 files changed, 112 insertions(+), 158 deletions(-) rename group-coordinator/src/test/java/org/apache/kafka/coordinator/group/{consumer => }/AssignmentTestUtil.java (97%) diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTestUtil.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/AssignmentTestUtil.java similarity index 97% rename from group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTestUtil.java rename to group-coordinator/src/test/java/org/apache/kafka/coordinator/group/AssignmentTestUtil.java index 8b5e072aaa05e..acf0ae141e62a 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTestUtil.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/AssignmentTestUtil.java @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.kafka.coordinator.group.consumer; +package org.apache.kafka.coordinator.group; import org.apache.kafka.common.Uuid; diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 86a13fced762f..a1eda44e0c949 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -21,7 +21,6 @@ import org.junit.jupiter.api.Test; import java.util.Arrays; -import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -30,6 +29,8 @@ import java.util.Set; import java.util.TreeMap; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -93,26 +94,26 @@ public void testFirstAssignmentTwoConsumersTwoTopicsSameSubscriptions() { Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic3Uuid), - Collections.emptyMap()) - ); + Collections.emptyMap() + )); members.put(consumerB, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic3Uuid), - Collections.emptyMap()) - ); + Collections.emptyMap() + )); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - Map>> expectedAssignment = new HashMap<>(); + Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1)); - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(2)); + expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1))); + expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 2))); // Topic 3 Partitions Assignment - mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(0)); - mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(1)); + expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 0))); + expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 1))); assertAssignment(expectedAssignment, computedAssignment); } @@ -130,35 +131,35 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - Collections.emptyMap()) - ); + Collections.emptyMap() + )); members.put(consumerB, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), Collections.singletonList(topic3Uuid), - Collections.emptyMap()) - ); + Collections.emptyMap() + )); members.put(consumerC, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), Arrays.asList(topic2Uuid, topic3Uuid), - Collections.emptyMap()) - ); + Collections.emptyMap() + )); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - Map>> expectedAssignment = new HashMap<>(); + Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1, 2)); + expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1, 2))); // Topic 2 Partitions Assignment - mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(0, 1)); - mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(2)); + expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0, 1))); + expectedAssignment.put(consumerC, mkAssignment(mkTopicAssignment(topic2Uuid, 2))); // Topic 3 Partitions Assignment - mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(0)); - mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(1)); + expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic3Uuid, 0))); + expectedAssignment.get(consumerC).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 1))); assertAssignment(expectedAssignment, computedAssignment); } @@ -175,35 +176,35 @@ public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic3Uuid), - Collections.emptyMap()) - ); + Collections.emptyMap() + )); members.put(consumerB, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic3Uuid), - Collections.emptyMap()) - ); + Collections.emptyMap() + )); members.put(consumerC, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic3Uuid), - Collections.emptyMap()) - ); + Collections.emptyMap() + )); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - Map>> expectedAssignment = new HashMap<>(); + Map>> expectedAssignment = new HashMap<>(); // Topic 3 has 2 partitions but three consumers subscribed to it - one of them will not get a partition. // Topic 1 Partitions Assignment - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(0)); - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(1)); - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(2)); - // Topic 2 Partitions Assignment - mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(0)); - mkAssignment(expectedAssignment, topic3Uuid, Collections.singleton(1)); + expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0))); + expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 1))); + expectedAssignment.put(consumerC, mkAssignment(mkTopicAssignment(topic1Uuid, 2))); + // Topic 3 Partitions Assignment + expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 0))); + expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 1))); assertAssignment(expectedAssignment, computedAssignment); } @@ -223,8 +224,8 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForA) - ); + currentAssignmentForA + )); Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic1Uuid, Collections.singleton(1)); @@ -233,35 +234,27 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForB) - ); + currentAssignmentForB + )); // Add a new consumer to trigger a re-assignment members.put(consumerC, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - Collections.emptyMap()) - ); + Collections.emptyMap() + )); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - Map>> expectedAssignment = new HashMap<>(); + Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(0)); - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(1)); + expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0))); + expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 1))); // Topic 2 Partitions Assignment - mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(0)); - mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(1)); - - // Test for stickiness - assertEquals(computedAssignment.members().get(consumerA) - .targetPartitions(), new HashMap<>(currentAssignmentForA), - "Stickiness test failed for consumer A"); - assertEquals(computedAssignment.members().get(consumerB) - .targetPartitions(), new HashMap<>(currentAssignmentForB), - "Stickiness test failed for consumer B"); + expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0))); + expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 1))); // Consumer C shouldn't get any assignment, due to stickiness A, B retain their assignments assertNull(computedAssignment.members().get(consumerC)); @@ -285,8 +278,8 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForA) - ); + currentAssignmentForA + )); Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic1Uuid, Collections.singleton(2)); @@ -295,25 +288,19 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForB) - ); + currentAssignmentForB + )); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - Map>> expectedAssignment = new HashMap<>(); + Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1)); - mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(2, 3)); + expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1))); + expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 2, 3))); // Topic 2 Partitions Assignment - mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(0, 1)); - mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(2, 3)); - - // Test for stickiness - assertEquals(computedAssignment.members().get(consumerA) - .targetPartitions(), new HashMap<>(currentAssignmentForA), - "Stickiness test failed for consumer A"); - // Implicitly stickiness is checked for B also since the other set of assigned partitions is (2,3) + expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0, 1))); + expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 2, 3))); assertAssignment(expectedAssignment, computedAssignment); assertCoPartitionJoinProperty(computedAssignment); @@ -325,7 +312,7 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); - Map members = new HashMap<>(); + Map members = new TreeMap<>(); Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); @@ -334,8 +321,8 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForA) - ); + currentAssignmentForA + )); Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic1Uuid, Collections.singleton(2)); @@ -344,40 +331,34 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForB) - ); + currentAssignmentForB + )); // Add a new consumer to trigger a re-assignment members.put(consumerC, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - Collections.emptyMap())); + Collections.emptyMap() + )); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - Map>> expectedAssignment = new HashMap<>(); + Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(0)); - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(2)); - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(1)); + expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0))); + expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 2))); + expectedAssignment.put(consumerC, mkAssignment(mkTopicAssignment(topic1Uuid, 1))); // Topic 2 Partitions Assignment - mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(0)); - mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(2)); - mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(1)); - - // Test for stickiness - assertEquals(computedAssignment.members().get(consumerB) - .targetPartitions(), new HashMap<>(currentAssignmentForB), - "Stickiness test failed for Consumer B"); - assertTrue(computedAssignment.members().get(consumerA) - .targetPartitions().get(topic1Uuid).contains(0), - "Stickiness test failed for Consumer A"); + expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0))); + expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 2))); + expectedAssignment.get(consumerC).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 1))); assertAssignment(expectedAssignment, computedAssignment); assertCoPartitionJoinProperty(computedAssignment); } + @Test public void testReassignmentWhenOneConsumerAddedAndOnePartitionAfterInitialAssignmentWithTwoConsumersTwoTopics() { Map topics = new HashMap<>(); @@ -385,7 +366,7 @@ public void testReassignmentWhenOneConsumerAddedAndOnePartitionAfterInitialAssig topics.put(topic1Uuid, new AssignmentTopicMetadata(4)); topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); - Map members = new HashMap<>(); + Map members = new TreeMap<>(); Map> currentAssignmentForA = new HashMap<>(); currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); @@ -394,8 +375,8 @@ public void testReassignmentWhenOneConsumerAddedAndOnePartitionAfterInitialAssig Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForA) - ); + currentAssignmentForA + )); Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic1Uuid, Collections.singleton(2)); @@ -404,36 +385,29 @@ public void testReassignmentWhenOneConsumerAddedAndOnePartitionAfterInitialAssig Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForB) - ); + currentAssignmentForB + )); // Add a new consumer to trigger a re-assignment members.put(consumerC, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), Collections.singletonList(topic1Uuid), - Collections.emptyMap())); + Collections.emptyMap() + )); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - Map>> expectedAssignment = new HashMap<>(); + Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1)); - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(2)); - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(3)); + expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1))); + expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 2))); + expectedAssignment.put(consumerC, mkAssignment(mkTopicAssignment(topic1Uuid, 3))); // Topic 2 Partitions Assignment // Since the new consumer isn't subscribed to topic 2 the assignment shouldn't change - mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(0, 1)); - mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(2)); - System.out.println("assignment is " + computedAssignment); - // Test for stickiness - assertEquals(computedAssignment.members().get(consumerA) - .targetPartitions(), new HashMap<>(currentAssignmentForA), - "Stickiness test failed for Consumer A"); - assertEquals(computedAssignment.members().get(consumerB) - .targetPartitions(), new HashMap<>(currentAssignmentForB), - "Stickiness test failed for Consumer B"); + expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0, 1))); + expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 2))); assertAssignment(expectedAssignment, computedAssignment); } @@ -444,7 +418,7 @@ public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoC topics.put(topic1Uuid, new AssignmentTopicMetadata(3)); topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); - Map members = new HashMap<>(); + Map members = new TreeMap<>(); // Consumer A was removed Map> currentAssignmentForB = new HashMap<>(); @@ -454,17 +428,17 @@ public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoC Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid), - currentAssignmentForB) - ); + currentAssignmentForB + )); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - Map>> expectedAssignment = new HashMap<>(); + Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1, 2)); + expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1, 2))); // Topic 2 Partitions Assignment - mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(0, 1, 2)); + expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0, 1, 2))); assertAssignment(expectedAssignment, computedAssignment); assertCoPartitionJoinProperty(computedAssignment); @@ -477,7 +451,7 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme topics.put(topic2Uuid, new AssignmentTopicMetadata(3)); topics.put(topic3Uuid, new AssignmentTopicMetadata(2)); - Map members = new HashMap<>(); + Map members = new TreeMap<>(); // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3 // Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2 @@ -489,8 +463,8 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme Optional.empty(), Optional.empty(), Collections.singletonList(topic1Uuid), - currentAssignmentForA) - ); + currentAssignmentForA + )); Map> currentAssignmentForB = new HashMap<>(); currentAssignmentForB.put(topic2Uuid, Collections.singleton(1)); @@ -498,8 +472,8 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme Optional.empty(), Optional.empty(), Arrays.asList(topic1Uuid, topic2Uuid, topic3Uuid), - currentAssignmentForB) - ); + currentAssignmentForB + )); Map> currentAssignmentForC = new HashMap<>(); currentAssignmentForC.put(topic2Uuid, Collections.singleton(2)); @@ -508,49 +482,29 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme Optional.empty(), Optional.empty(), Collections.singletonList(topic2Uuid), - currentAssignmentForC) - ); + currentAssignmentForC + )); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); GroupAssignment computedAssignment = assignor.assign(assignmentSpec); - Map>> expectedAssignment = new HashMap<>(); + Map>> expectedAssignment = new HashMap<>(); // Topic 1 Partitions Assignment - mkAssignment(expectedAssignment, topic1Uuid, Arrays.asList(0, 1)); - mkAssignment(expectedAssignment, topic1Uuid, Collections.singleton(2)); + expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1))); + expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 2))); // Topic 2 Partitions Assignment - mkAssignment(expectedAssignment, topic2Uuid, Arrays.asList(0, 1)); - mkAssignment(expectedAssignment, topic2Uuid, Collections.singleton(2)); + expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0, 1))); + expectedAssignment.put(consumerC, mkAssignment(mkTopicAssignment(topic2Uuid, 2))); // Topic 3 Partitions Assignment - mkAssignment(expectedAssignment, topic3Uuid, Arrays.asList(0, 1)); - - // Test for stickiness - assertTrue(computedAssignment.members().get(consumerC) - .targetPartitions().get(topic2Uuid).contains(2), - "Stickiness test failed for Consumer C"); - assertTrue(computedAssignment.members().get(consumerA) - .targetPartitions().get(topic1Uuid).containsAll(Arrays.asList(0, 1)), - "Stickiness test failed for Consumer A"); + expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 0, 1))); assertAssignment(expectedAssignment, computedAssignment); } - private void mkAssignment(Map>> expectedAssignment, Uuid topicId, Collection partitions) { - expectedAssignment.computeIfAbsent(topicId, k -> new HashSet<>()).add(new HashSet<>(partitions)); - } - - // We have a set of sets with the partitions that should be distributed amongst the consumers, if it exists then remove it from the set. - // The test is done like this since the order in which members are assigned partitions isn't guaranteed. We are just testing if the computed - // assignment contains the expected set of partitions, irrespective of which member got them. - private void assertAssignment(Map>> expectedAssignment, GroupAssignment computedGroupAssignment) { - for (MemberAssignment member : computedGroupAssignment.members().values()) { - Map> computedAssignmentForMember = member.targetPartitions(); - for (Map.Entry> assignmentForTopic : computedAssignmentForMember.entrySet()) { - Uuid topicId = assignmentForTopic.getKey(); - Set assignmentPartitionsSet = assignmentForTopic.getValue(); - assertTrue(expectedAssignment.get(topicId).contains(assignmentPartitionsSet)); - expectedAssignment.remove(assignmentPartitionsSet); - } + private void assertAssignment(Map>> expectedAssignment, GroupAssignment computedGroupAssignment) { + for (String memberId : computedGroupAssignment.members().keySet()) { + Map> computedAssignmentForMember = computedGroupAssignment.members().get(memberId).targetPartitions(); + assertEquals(expectedAssignment.get(memberId), computedAssignmentForMember); } } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTest.java index d81006cc026e0..2cff004ded22f 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTest.java @@ -29,8 +29,8 @@ import java.util.Map; import java.util.Set; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkAssignment; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkTopicAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertThrows; diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java index 7df34650be87f..e98a895d2beb8 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java @@ -27,8 +27,8 @@ import java.util.Optional; import java.util.OptionalInt; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkAssignment; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkTopicAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment; import static org.junit.jupiter.api.Assertions.assertEquals; public class ConsumerGroupMemberTest { From 56863d886af90fe1bd6a3533adff4926237b9db1 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 9 May 2023 01:12:07 -0700 Subject: [PATCH 36/44] not using code for variable names in java doc --- .../coordinator/group/assignor/RangeAssignor.java | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index a07aa30d5703f..b4008cec05151 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -104,13 +104,13 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS /** *

    The algorithm includes the following steps: *

      - *
    1. Generate a map of membersPerTopic using the given member subscriptions.
      2. - *
    2. Generate a list of members (potentiallyUnfilledMembers) that have not met the minimum required quota of partitions for the assignment AND - * get a list (assignedStickyPartitionsPerTopic) of partitions that will be retained in the new assignment.
      3. - *
    3. Add members from the potentiallyUnfilledMembers list to the unfilledMembersPerTopic map if they haven't met the total required quota + *
    4. Generate a map of members per topic using the given member subscriptions.
      5. + *
    5. Generate a list of members called potentially unfilled members, which consists of members that have not met the minimum required quota of partitions for the assignment AND + * get a list called assigned sticky partitions per topic, which has the partitions that will be retained in the new assignment.
      6. + *
    6. Add members from the potentially unfilled members list to the unfilled members per topic map if they haven't met the total required quota * i.e. minRequiredQuota + 1, if the member is designated to receive one of the excess partitions OR minRequiredQuota otherwise.
      7. *
    7. Generate a list of unassigned partitions by calculating the difference between the total partitions for the topic and the assigned (sticky) partitions.
      8. - *
    8. Check if unfilled members exist for the current topicId and assign partitions to them in ranges from the unassignedPartitionsPerTopic map + *
    9. Check if unfilled members exist for the current topicId and assign partitions to them in ranges from the unassigned partitions per topic map * based on the remaining partitions value stored.
      10. *
    *

    From 4a3ad5faad6c0f899231aad58288abbc6e142746 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 9 May 2023 01:40:01 -0700 Subject: [PATCH 37/44] changed code --- .../group/assignor/RangeAssignor.java | 59 ++++++++----------- 1 file changed, 25 insertions(+), 34 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index b4008cec05151..c501b4ab1fe45 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -107,11 +107,11 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS *
  • Generate a map of members per topic using the given member subscriptions.
    • *
  • Generate a list of members called potentially unfilled members, which consists of members that have not met the minimum required quota of partitions for the assignment AND * get a list called assigned sticky partitions per topic, which has the partitions that will be retained in the new assignment.
    • - *
  • Add members from the potentially unfilled members list to the unfilled members per topic map if they haven't met the total required quota - * i.e. minRequiredQuota + 1, if the member is designated to receive one of the excess partitions OR minRequiredQuota otherwise.
    • *
  • Generate a list of unassigned partitions by calculating the difference between the total partitions for the topic and the assigned (sticky) partitions.
    • - *
  • Check if unfilled members exist for the current topicId and assign partitions to them in ranges from the unassigned partitions per topic map - * based on the remaining partitions value stored.
    • + *
  • Find members from the potentially unfilled members list that haven't met the total required quota + * i.e. minRequiredQuota + 1, if the member is designated to receive one of the excess partitions OR minRequiredQuota otherwise.
    • + *
  • Assign partitions to them in ranges from the unassigned partitions per topic + * based on the remaining partitions value.
    • * *

    */ @@ -188,25 +188,7 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit } } - List unfilledMembersForTopic = new ArrayList<>(); - // Step 3 - // If remaining > 0 after increasing the required quota due to the extra partition, add potentially unfilled member to the unfilled members list. - for (MemberWithRemainingAssignments pair : potentiallyUnfilledMembers) { - String memberId = pair.memberId; - Integer remaining = pair.remaining; - if (numMembersWithExtraPartition > 0) { - remaining++; - numMembersWithExtraPartition--; - } - if (remaining > 0) { - MemberWithRemainingAssignments newPair = new MemberWithRemainingAssignments(memberId, remaining); - unfilledMembersForTopic - .add(newPair); - } - } - - // Step 4 // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get unassigned partitions. // List of unassigned partitions per topic contains the partitions in ascending order. List unassignedPartitionsForTopic = new ArrayList<>(); @@ -216,19 +198,28 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit } } - // Step 5 - // Assign the remaining number of partitions from the unassigned partitions list. + List unfilledMembersForTopic = new ArrayList<>(); + // Step 4 + // If remaining > 0 after increasing the required quota due to the extra partition, add potentially unfilled member to the unfilled members list. int unassignedPartitionsListStartPointer = 0; - for (MemberWithRemainingAssignments memberWithRemainingAssignments : unfilledMembersForTopic) { - String memberId = memberWithRemainingAssignments.memberId; - int remaining = memberWithRemainingAssignments.remaining; - List partitionsToAssign = unassignedPartitionsForTopic - .subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); - unassignedPartitionsListStartPointer += remaining; - newAssignment.computeIfAbsent(memberId, k -> new MemberAssignment(new HashMap<>())) - .targetPartitions() - .computeIfAbsent(topicId, k -> new HashSet<>()) - .addAll(partitionsToAssign); + for (MemberWithRemainingAssignments pair : potentiallyUnfilledMembers) { + String memberId = pair.memberId; + Integer remaining = pair.remaining; + if (numMembersWithExtraPartition > 0) { + remaining++; + numMembersWithExtraPartition--; + } + if (remaining > 0) { + // Step 5 + // Assign the remaining number of partitions from the unassigned partitions list. + List partitionsToAssign = unassignedPartitionsForTopic + .subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); + unassignedPartitionsListStartPointer += remaining; + newAssignment.computeIfAbsent(memberId, k -> new MemberAssignment(new HashMap<>())) + .targetPartitions() + .computeIfAbsent(topicId, k -> new HashSet<>()) + .addAll(partitionsToAssign); + } } }); From 85766ce162da453049dc78f25a2b5bea5b77340d Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 9 May 2023 01:43:20 -0700 Subject: [PATCH 38/44] minor --- .../kafka/coordinator/group/assignor/RangeAssignor.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index c501b4ab1fe45..960c0fca8a40c 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -198,7 +198,6 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit } } - List unfilledMembersForTopic = new ArrayList<>(); // Step 4 // If remaining > 0 after increasing the required quota due to the extra partition, add potentially unfilled member to the unfilled members list. int unassignedPartitionsListStartPointer = 0; @@ -218,7 +217,8 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit newAssignment.computeIfAbsent(memberId, k -> new MemberAssignment(new HashMap<>())) .targetPartitions() .computeIfAbsent(topicId, k -> new HashSet<>()) - .addAll(partitionsToAssign); + .addAll(partitionsToAssign + ); } } }); From 9bcaf34d417404e14e2407107534f5ef2073dbb5 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 9 May 2023 01:50:12 -0700 Subject: [PATCH 39/44] minor --- .../apache/kafka/coordinator/group/assignor/RangeAssignor.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 960c0fca8a40c..03c59bde6a007 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -217,8 +217,7 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit newAssignment.computeIfAbsent(memberId, k -> new MemberAssignment(new HashMap<>())) .targetPartitions() .computeIfAbsent(topicId, k -> new HashSet<>()) - .addAll(partitionsToAssign - ); + .addAll(partitionsToAssign); } } }); From edbf2275dad4cf467c1144a918d141f89ec938f6 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 9 May 2023 16:38:33 -0700 Subject: [PATCH 40/44] html formatting changes and other small things --- .../group/assignor/RangeAssignor.java | 77 +++---- .../group/assignor/RangeAssignorTest.java | 199 ++++++++++-------- 2 files changed, 149 insertions(+), 127 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 03c59bde6a007..ffefaa5b28d79 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -35,17 +35,17 @@ * This Range Assignor inherits properties of both the range assignor and the sticky assignor. * The properties are as follows: *
      - *
    1. Each member must get at least one partition from every topic that it is subscribed to. The only exception is when - * the number of subscribed members is greater than the number of partitions for that topic. (Range)
      2. - *
    2. Partitions should be assigned to members in a way that facilitates the join operation when required. (Range) - * This can only be done if every member is subscribed to the same topics and the topics are co-partitioned. - * Two streams are co-partitioned if the following conditions are met: - *
        - *
      • The keys must have the same schemas.
        • - *
      • The topics involved must have the same number of partitions.
        • - *
      - *
      3. - *
    3. Members should retain as much of their previous assignment as possible to reduce the number of partition movements during reassignment. (Sticky)
      4. + *
    4. Each member must get at least one partition from every topic that it is subscribed to. The only exception is when + * the number of subscribed members is greater than the number of partitions for that topic. (Range)
      5. + *
    5. Partitions should be assigned to members in a way that facilitates the join operation when required. (Range) + * This can only be done if every member is subscribed to the same topics and the topics are co-partitioned. + * Two streams are co-partitioned if the following conditions are met: + *
        + *
      • The keys must have the same schemas.
        • + *
      • The topics involved must have the same number of partitions.
        • + *
      + *
      6. + *
    6. Members should retain as much of their previous assignment as possible to reduce the number of partition movements during reassignment. (Sticky)
      7. *
    */ public class RangeAssignor implements PartitionAssignor { @@ -78,7 +78,7 @@ public MemberWithRemainingAssignments(String memberId, Integer remaining) { } /** - * @return Map of topicIds to a list of members subscribed to them. + * @return Map of topic ids to a list of members subscribed to them. */ private Map> membersPerTopic(final AssignmentSpec assignmentSpec) { Map> membersPerTopic = new HashMap<>(); @@ -86,7 +86,7 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS membersData.forEach((memberId, memberMetadata) -> { Collection topics = memberMetadata.subscribedTopicIds(); - for (Uuid topicId: topics) { + for (Uuid topicId : topics) { // Only topics that are present in both the subscribed topics list and the topic metadata should be considered for assignment. if (assignmentSpec.topics().containsKey(topicId)) { membersPerTopic @@ -104,14 +104,14 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS /** *

    The algorithm includes the following steps: *

      - *
    1. Generate a map of members per topic using the given member subscriptions.
      2. - *
    2. Generate a list of members called potentially unfilled members, which consists of members that have not met the minimum required quota of partitions for the assignment AND - * get a list called assigned sticky partitions per topic, which has the partitions that will be retained in the new assignment.
      3. - *
    3. Generate a list of unassigned partitions by calculating the difference between the total partitions for the topic and the assigned (sticky) partitions.
      4. - *
    4. Find members from the potentially unfilled members list that haven't met the total required quota - * i.e. minRequiredQuota + 1, if the member is designated to receive one of the excess partitions OR minRequiredQuota otherwise.
      5. - *
    5. Assign partitions to them in ranges from the unassigned partitions per topic - * based on the remaining partitions value.
      6. + *
    6. Generate a map of members per topic using the given member subscriptions.
      7. + *
    7. Generate a list of members called potentially unfilled members, which consists of members that have not met the minimum required quota of partitions for the assignment AND + * get a list called assigned sticky partitions for topic, which has the partitions that will be retained in the new assignment.
      8. + *
    8. Generate a list of unassigned partitions by calculating the difference between the total partitions for the topic and the assigned (sticky) partitions.
      9. + *
    9. Find members from the potentially unfilled members list that haven't met the total required quota + * i.e. minRequiredQuota + 1, if the member is designated to receive one of the excess partitions OR minRequiredQuota otherwise.
      10. + *
    10. Assign partitions to them in ranges from the unassigned partitions per topic + * based on the remaining partitions value.
      11. *
    *

    */ @@ -130,9 +130,6 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit // Each member can get only ONE extra partition per topic after receiving the minimum quota. int numMembersWithExtraPartition = numPartitionsForTopic % membersForTopic.size(); - // Idle members case : When the number of members subscribed to a topic is greater than the total number of partitions, - // all members get assigned via the "extra partitions" logic since minRequiredQuota = 0. - // Step 2 Set assignedStickyPartitionsForTopic = new HashSet<>(); List potentiallyUnfilledMembers = new ArrayList<>(); @@ -144,14 +141,11 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit int currentAssignmentSize = assignedPartitionsForTopic.size(); List currentAssignmentListForTopic = new ArrayList<>(assignedPartitionsForTopic); - // If there were partitions from this topic that were previously assigned to this member present, retain as much as possible. - // We need to retain: - // 1) currentSize number of partitions (when currentSize < required) OR - // 2) required number of partitions (when currentSize > required). + // If there were partitions from this topic that were previously assigned to this member, retain as many as possible. + // Sort the current assignment in ascending order since we want the same partition numbers from each topic + // to go to the same member, in order to facilitate joins in case of co-partitioned topics. if (currentAssignmentSize > 0) { int retainedPartitionsCount = min(currentAssignmentSize, minRequiredQuota); - // The partitions need to be in ascending order since we want the same partition numbers from each topic - // to go to the same member to facilitate joins in case of co-partitioned topics. Collections.sort(currentAssignmentListForTopic); for (int i = 0; i < retainedPartitionsCount; i++) { assignedStickyPartitionsForTopic @@ -170,8 +164,7 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit // 3) If remaining > 0: member doesn't have the minimum required partitions, so it should be added to potentiallyUnfilledMembers. int remaining = minRequiredQuota - currentAssignmentSize; - // In order to remain as sticky as possible, since the order of members can be different, we want the members that have more than the required number of partitions, - // to get the extra partition instead of assigning the extras to the first few members directly. + // Retain extra partitions as well when applicable. if (remaining < 0 && numMembersWithExtraPartition > 0) { numMembersWithExtraPartition--; // Since we already added the minimumRequiredQuota of partitions in the previous step (until minReq - 1), we just need to @@ -189,8 +182,8 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit } // Step 3 - // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get unassigned partitions. - // List of unassigned partitions per topic contains the partitions in ascending order. + // Find the difference between the total partitions per topic and the already assigned sticky partitions for the topic to get the unassigned partitions. + // List of unassigned partitions for topic contains the partitions in ascending order. List unassignedPartitionsForTopic = new ArrayList<>(); for (int i = 0; i < numPartitionsForTopic; i++) { if (!assignedStickyPartitionsForTopic.contains(i)) { @@ -198,8 +191,9 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit } } - // Step 4 - // If remaining > 0 after increasing the required quota due to the extra partition, add potentially unfilled member to the unfilled members list. + // Step 4 and Step 5 + // Account for the extra partitions if necessary and increase the required quota by 1. + // If remaining > 0 after increasing the required quota, assign the remaining number of partitions from the unassigned partitions list. int unassignedPartitionsListStartPointer = 0; for (MemberWithRemainingAssignments pair : potentiallyUnfilledMembers) { String memberId = pair.memberId; @@ -209,19 +203,16 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit numMembersWithExtraPartition--; } if (remaining > 0) { - // Step 5 - // Assign the remaining number of partitions from the unassigned partitions list. List partitionsToAssign = unassignedPartitionsForTopic - .subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); + .subList(unassignedPartitionsListStartPointer, unassignedPartitionsListStartPointer + remaining); unassignedPartitionsListStartPointer += remaining; newAssignment.computeIfAbsent(memberId, k -> new MemberAssignment(new HashMap<>())) - .targetPartitions() - .computeIfAbsent(topicId, k -> new HashSet<>()) - .addAll(partitionsToAssign); + .targetPartitions() + .computeIfAbsent(topicId, k -> new HashSet<>()) + .addAll(partitionsToAssign); } } }); - return new GroupAssignment(newAssignment); } } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index a1eda44e0c949..8198e02dea8fe 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -108,12 +108,16 @@ public void testFirstAssignmentTwoConsumersTwoTopicsSameSubscriptions() { GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1))); - expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 2))); - // Topic 3 Partitions Assignment - expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 0))); - expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 1))); + + expectedAssignment.put(consumerA, mkAssignment( + mkTopicAssignment(topic1Uuid, 0, 1), + mkTopicAssignment(topic3Uuid, 0) + )); + + expectedAssignment.put(consumerB, mkAssignment( + mkTopicAssignment(topic1Uuid, 2), + mkTopicAssignment(topic3Uuid, 1) + )); assertAssignment(expectedAssignment, computedAssignment); } @@ -152,14 +156,20 @@ public void testFirstAssignmentThreeConsumersThreeTopicsDifferentSubscriptions() GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1, 2))); - // Topic 2 Partitions Assignment - expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0, 1))); - expectedAssignment.put(consumerC, mkAssignment(mkTopicAssignment(topic2Uuid, 2))); - // Topic 3 Partitions Assignment - expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic3Uuid, 0))); - expectedAssignment.get(consumerC).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 1))); + + expectedAssignment.put(consumerA, mkAssignment( + mkTopicAssignment(topic1Uuid, 0, 1, 2), + mkTopicAssignment(topic2Uuid, 0, 1) + )); + + expectedAssignment.put(consumerB, mkAssignment( + mkTopicAssignment(topic3Uuid, 0) + )); + + expectedAssignment.put(consumerC, mkAssignment( + mkTopicAssignment(topic2Uuid, 2), + mkTopicAssignment(topic3Uuid, 1) + )); assertAssignment(expectedAssignment, computedAssignment); } @@ -198,13 +208,19 @@ public void testFirstAssignmentNumConsumersGreaterThanNumPartitions() { Map>> expectedAssignment = new HashMap<>(); // Topic 3 has 2 partitions but three consumers subscribed to it - one of them will not get a partition. - // Topic 1 Partitions Assignment - expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0))); - expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 1))); - expectedAssignment.put(consumerC, mkAssignment(mkTopicAssignment(topic1Uuid, 2))); - // Topic 3 Partitions Assignment - expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 0))); - expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 1))); + expectedAssignment.put(consumerA, mkAssignment( + mkTopicAssignment(topic1Uuid, 0), + mkTopicAssignment(topic3Uuid, 0) + )); + + expectedAssignment.put(consumerB, mkAssignment( + mkTopicAssignment(topic1Uuid, 1), + mkTopicAssignment(topic3Uuid, 1) + )); + + expectedAssignment.put(consumerC, mkAssignment( + mkTopicAssignment(topic1Uuid, 2) + )); assertAssignment(expectedAssignment, computedAssignment); } @@ -217,9 +233,11 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA Map members = new TreeMap<>(); - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, Collections.singleton(0)); - currentAssignmentForA.put(topic2Uuid, Collections.singleton(0)); + Map> currentAssignmentForA = mkAssignment( + mkTopicAssignment(topic1Uuid, 0), + mkTopicAssignment(topic2Uuid, 0) + ); + members.put(consumerA, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -227,9 +245,11 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA currentAssignmentForA )); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, Collections.singleton(1)); - currentAssignmentForB.put(topic2Uuid, Collections.singleton(1)); + Map> currentAssignmentForB = mkAssignment( + mkTopicAssignment(topic1Uuid, 1), + mkTopicAssignment(topic2Uuid, 1) + ); + members.put(consumerB, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -249,17 +269,20 @@ public void testReassignmentNumConsumersGreaterThanNumPartitionsWhenOneConsumerA GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0))); - expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 1))); - // Topic 2 Partitions Assignment - expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0))); - expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 1))); + + expectedAssignment.put(consumerA, mkAssignment( + mkTopicAssignment(topic1Uuid, 0), + mkTopicAssignment(topic2Uuid, 0) + )); + + expectedAssignment.put(consumerB, mkAssignment( + mkTopicAssignment(topic1Uuid, 1), + mkTopicAssignment(topic2Uuid, 1) + )); // Consumer C shouldn't get any assignment, due to stickiness A, B retain their assignments assertNull(computedAssignment.members().get(consumerC)); assertAssignment(expectedAssignment, computedAssignment); - assertCoPartitionJoinProperty(computedAssignment); } @Test @@ -295,15 +318,18 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1))); - expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 2, 3))); - // Topic 2 Partitions Assignment - expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0, 1))); - expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 2, 3))); + + expectedAssignment.put(consumerA, mkAssignment( + mkTopicAssignment(topic1Uuid, 0, 1), + mkTopicAssignment(topic2Uuid, 0, 1) + )); + + expectedAssignment.put(consumerB, mkAssignment( + mkTopicAssignment(topic1Uuid, 2, 3), + mkTopicAssignment(topic2Uuid, 2, 3) + )); assertAssignment(expectedAssignment, computedAssignment); - assertCoPartitionJoinProperty(computedAssignment); } @Test @@ -346,17 +372,23 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0))); - expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 2))); - expectedAssignment.put(consumerC, mkAssignment(mkTopicAssignment(topic1Uuid, 1))); - // Topic 2 Partitions Assignment - expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0))); - expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 2))); - expectedAssignment.get(consumerC).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 1))); + + expectedAssignment.put(consumerA, mkAssignment( + mkTopicAssignment(topic1Uuid, 0), + mkTopicAssignment(topic2Uuid, 0) + )); + + expectedAssignment.put(consumerB, mkAssignment( + mkTopicAssignment(topic1Uuid, 2), + mkTopicAssignment(topic2Uuid, 2) + )); + + expectedAssignment.put(consumerC, mkAssignment( + mkTopicAssignment(topic1Uuid, 1), + mkTopicAssignment(topic2Uuid, 1) + )); assertAssignment(expectedAssignment, computedAssignment); - assertCoPartitionJoinProperty(computedAssignment); } @Test @@ -400,14 +432,20 @@ public void testReassignmentWhenOneConsumerAddedAndOnePartitionAfterInitialAssig GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1))); - expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 2))); - expectedAssignment.put(consumerC, mkAssignment(mkTopicAssignment(topic1Uuid, 3))); - // Topic 2 Partitions Assignment - // Since the new consumer isn't subscribed to topic 2 the assignment shouldn't change - expectedAssignment.get(consumerA).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0, 1))); - expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 2))); + + expectedAssignment.put(consumerA, mkAssignment( + mkTopicAssignment(topic1Uuid, 0, 1), + mkTopicAssignment(topic2Uuid, 0, 1) + )); + + expectedAssignment.put(consumerB, mkAssignment( + mkTopicAssignment(topic1Uuid, 2), + mkTopicAssignment(topic2Uuid, 2) + )); + + expectedAssignment.put(consumerC, mkAssignment( + mkTopicAssignment(topic1Uuid, 3) + )); assertAssignment(expectedAssignment, computedAssignment); } @@ -435,13 +473,13 @@ public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoC GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1, 2))); - // Topic 2 Partitions Assignment - expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0, 1, 2))); + + expectedAssignment.put(consumerB, mkAssignment( + mkTopicAssignment(topic1Uuid, 0, 1, 2), + mkTopicAssignment(topic2Uuid, 0, 1, 2) + )); assertAssignment(expectedAssignment, computedAssignment); - assertCoPartitionJoinProperty(computedAssignment); } @Test @@ -489,14 +527,20 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme GroupAssignment computedAssignment = assignor.assign(assignmentSpec); Map>> expectedAssignment = new HashMap<>(); - // Topic 1 Partitions Assignment - expectedAssignment.put(consumerA, mkAssignment(mkTopicAssignment(topic1Uuid, 0, 1))); - expectedAssignment.put(consumerB, mkAssignment(mkTopicAssignment(topic1Uuid, 2))); - // Topic 2 Partitions Assignment - expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic2Uuid, 0, 1))); - expectedAssignment.put(consumerC, mkAssignment(mkTopicAssignment(topic2Uuid, 2))); - // Topic 3 Partitions Assignment - expectedAssignment.get(consumerB).putAll(mkAssignment(mkTopicAssignment(topic3Uuid, 0, 1))); + + expectedAssignment.put(consumerA, mkAssignment( + mkTopicAssignment(topic1Uuid, 0, 1) + )); + + expectedAssignment.put(consumerB, mkAssignment( + mkTopicAssignment(topic1Uuid, 2), + mkTopicAssignment(topic2Uuid, 0, 1), + mkTopicAssignment(topic3Uuid, 0, 1) + )); + + expectedAssignment.put(consumerC, mkAssignment( + mkTopicAssignment(topic2Uuid, 2) + )); assertAssignment(expectedAssignment, computedAssignment); } @@ -507,17 +551,4 @@ private void assertAssignment(Map>> expectedAssig assertEquals(expectedAssignment.get(memberId), computedAssignmentForMember); } } - - private void assertCoPartitionJoinProperty(GroupAssignment groupAssignment) { - for (MemberAssignment member : groupAssignment.members().values()) { - Map> computedAssignmentForMember = member.targetPartitions(); - Set compareSet = new HashSet<>(); - for (Set partitionsForTopicSet : computedAssignmentForMember.values()) { - if (compareSet.isEmpty()) { - compareSet = partitionsForTopicSet; - } - assertEquals(compareSet, partitionsForTopicSet); - } - } - } } From 62470e8462da0bcee568803fc9ba7d7efc046716 Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 9 May 2023 16:54:41 -0700 Subject: [PATCH 41/44] Test formatting changes --- .../group/assignor/RangeAssignorTest.java | 79 ++++++++++++------- 1 file changed, 49 insertions(+), 30 deletions(-) diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 8198e02dea8fe..cfeee6c8e407e 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -23,7 +23,6 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; -import java.util.HashSet; import java.util.Map; import java.util.Optional; import java.util.Set; @@ -294,9 +293,11 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { Map members = new TreeMap<>(); - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); + Map> currentAssignmentForA = mkAssignment( + mkTopicAssignment(topic1Uuid, 0, 1), + mkTopicAssignment(topic2Uuid, 0, 1) + ); + members.put(consumerA, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -304,9 +305,11 @@ public void testReassignmentWhenOnePartitionAddedForTwoConsumersTwoTopics() { currentAssignmentForA )); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, Collections.singleton(2)); - currentAssignmentForB.put(topic2Uuid, Collections.singleton(2)); + Map> currentAssignmentForB = mkAssignment( + mkTopicAssignment(topic1Uuid, 2), + mkTopicAssignment(topic2Uuid, 2) + ); + members.put(consumerB, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -340,9 +343,11 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon Map members = new TreeMap<>(); - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); + Map> currentAssignmentForA = mkAssignment( + mkTopicAssignment(topic1Uuid, 0, 1), + mkTopicAssignment(topic2Uuid, 0, 1) + ); + members.put(consumerA, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -350,9 +355,11 @@ public void testReassignmentWhenOneConsumerAddedAfterInitialAssignmentWithTwoCon currentAssignmentForA )); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, Collections.singleton(2)); - currentAssignmentForB.put(topic2Uuid, Collections.singleton(2)); + Map> currentAssignmentForB = mkAssignment( + mkTopicAssignment(topic1Uuid, 2), + mkTopicAssignment(topic2Uuid, 2) + ); + members.put(consumerB, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -400,9 +407,11 @@ public void testReassignmentWhenOneConsumerAddedAndOnePartitionAfterInitialAssig Map members = new TreeMap<>(); - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1))); - currentAssignmentForA.put(topic2Uuid, new HashSet<>(Arrays.asList(0, 1))); + Map> currentAssignmentForA = mkAssignment( + mkTopicAssignment(topic1Uuid, 0, 1), + mkTopicAssignment(topic2Uuid, 0, 1) + ); + members.put(consumerA, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -410,9 +419,11 @@ public void testReassignmentWhenOneConsumerAddedAndOnePartitionAfterInitialAssig currentAssignmentForA )); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, Collections.singleton(2)); - currentAssignmentForB.put(topic2Uuid, Collections.singleton(2)); + Map> currentAssignmentForB = mkAssignment( + mkTopicAssignment(topic1Uuid, 2), + mkTopicAssignment(topic2Uuid, 2) + ); + members.put(consumerB, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -459,9 +470,11 @@ public void testReassignmentWhenOneConsumerRemovedAfterInitialAssignmentWithTwoC Map members = new TreeMap<>(); // Consumer A was removed - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic1Uuid, new HashSet<>(Collections.singletonList(2))); - currentAssignmentForB.put(topic2Uuid, new HashSet<>(Collections.singletonList(2))); + Map> currentAssignmentForB = mkAssignment( + mkTopicAssignment(topic1Uuid, 2), + mkTopicAssignment(topic2Uuid, 2) + ); + members.put(consumerB, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -494,9 +507,11 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3 // Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2 - Map> currentAssignmentForA = new HashMap<>(); - currentAssignmentForA.put(topic1Uuid, new HashSet<>(Arrays.asList(0, 1, 2))); - currentAssignmentForA.put(topic2Uuid, Collections.singleton(0)); + Map> currentAssignmentForA = mkAssignment( + mkTopicAssignment(topic1Uuid, 0, 1, 2), + mkTopicAssignment(topic2Uuid, 0) + ); + members.put(consumerA, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -504,8 +519,10 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme currentAssignmentForA )); - Map> currentAssignmentForB = new HashMap<>(); - currentAssignmentForB.put(topic2Uuid, Collections.singleton(1)); + Map> currentAssignmentForB = mkAssignment( + mkTopicAssignment(topic2Uuid, 1) + ); + members.put(consumerB, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), @@ -513,9 +530,11 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme currentAssignmentForB )); - Map> currentAssignmentForC = new HashMap<>(); - currentAssignmentForC.put(topic2Uuid, Collections.singleton(2)); - currentAssignmentForC.put(topic3Uuid, new HashSet<>(Arrays.asList(0, 1))); + Map> currentAssignmentForC = mkAssignment( + mkTopicAssignment(topic2Uuid, 2), + mkTopicAssignment(topic3Uuid, 0, 1) + ); + members.put(consumerC, new AssignmentMemberSpec( Optional.empty(), Optional.empty(), From 9c9afd864c90e39fefa26f9cc5c8a77c745d9bff Mon Sep 17 00:00:00 2001 From: Ritika Reddy Date: Tue, 9 May 2023 17:18:06 -0700 Subject: [PATCH 42/44] addressed whatever was left --- .../kafka/coordinator/group/assignor/RangeAssignor.java | 4 +--- .../kafka/coordinator/group/assignor/RangeAssignorTest.java | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index ffefaa5b28d79..6c180dc276d20 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -123,9 +123,7 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit Map> membersPerTopic = membersPerTopic(assignmentSpec); membersPerTopic.forEach((topicId, membersForTopic) -> { - - AssignmentTopicMetadata topicData = assignmentSpec.topics().get(topicId); - int numPartitionsForTopic = topicData.numPartitions(); + int numPartitionsForTopic = assignmentSpec.topics().get(topicId).numPartitions(); int minRequiredQuota = numPartitionsForTopic / membersForTopic.size(); // Each member can get only ONE extra partition per topic after receiving the minimum quota. int numMembersWithExtraPartition = numPartitionsForTopic % membersForTopic.size(); diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index cfeee6c8e407e..8aff7e9190e1c 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -565,6 +565,7 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme } private void assertAssignment(Map>> expectedAssignment, GroupAssignment computedGroupAssignment) { + assertEquals(expectedAssignment.size(), computedGroupAssignment.members().size()); for (String memberId : computedGroupAssignment.members().keySet()) { Map> computedAssignmentForMember = computedGroupAssignment.members().get(memberId).targetPartitions(); assertEquals(expectedAssignment.get(memberId), computedAssignmentForMember); From 54ca849db5808a2a020c3779182f38b0fc96be51 Mon Sep 17 00:00:00 2001 From: David Jacot Date: Wed, 10 May 2023 08:54:55 +0200 Subject: [PATCH 43/44] small fixes and cosmetic changes --- .../group/assignor/RangeAssignor.java | 43 ++++++++++--------- .../group/assignor/RangeAssignorTest.java | 6 +-- 2 files changed, 26 insertions(+), 23 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java index 6c180dc276d20..b5316ffbbdbf2 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/RangeAssignor.java @@ -17,8 +17,6 @@ package org.apache.kafka.coordinator.group.assignor; import org.apache.kafka.common.Uuid; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import java.util.ArrayList; import java.util.Collection; @@ -36,21 +34,20 @@ * The properties are as follows: *
      *
    1. Each member must get at least one partition from every topic that it is subscribed to. The only exception is when - * the number of subscribed members is greater than the number of partitions for that topic. (Range)
      2. + * the number of subscribed members is greater than the number of partitions for that topic. (Range) *
    2. Partitions should be assigned to members in a way that facilitates the join operation when required. (Range) * This can only be done if every member is subscribed to the same topics and the topics are co-partitioned. * Two streams are co-partitioned if the following conditions are met: *
        - *
      • The keys must have the same schemas.
        • - *
      • The topics involved must have the same number of partitions.
        • + *
      • The keys must have the same schemas.
        • + *
      • The topics involved must have the same number of partitions.
        • *
      *
      3. - *
    3. Members should retain as much of their previous assignment as possible to reduce the number of partition movements during reassignment. (Sticky)
      4. + *
    4. Members should retain as much of their previous assignment as possible to reduce the number of partition + * movements during reassignment. (Sticky)
      5. *
    */ public class RangeAssignor implements PartitionAssignor { - private static final Logger log = LoggerFactory.getLogger(RangeAssignor.class); - public static final String RANGE_ASSIGNOR_NAME = "range"; @Override @@ -66,12 +63,13 @@ private static class MemberWithRemainingAssignments { * Member Id. */ private final String memberId; + /** * Number of partitions required to meet the assignment quota. */ - private final Integer remaining; + private final int remaining; - public MemberWithRemainingAssignments(String memberId, Integer remaining) { + public MemberWithRemainingAssignments(String memberId, int remaining) { this.memberId = memberId; this.remaining = remaining; } @@ -87,13 +85,15 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS membersData.forEach((memberId, memberMetadata) -> { Collection topics = memberMetadata.subscribedTopicIds(); for (Uuid topicId : topics) { - // Only topics that are present in both the subscribed topics list and the topic metadata should be considered for assignment. + // Only topics that are present in both the subscribed topics list and the topic metadata should be + // considered for assignment. if (assignmentSpec.topics().containsKey(topicId)) { membersPerTopic .computeIfAbsent(topicId, k -> new ArrayList<>()) .add(memberId); } else { - log.warn("Member " + memberId + " subscribed to topic " + topicId + " which doesn't exist in the topic metadata"); + throw new PartitionAssignorException("Member " + memberId + " subscribed to topic " + + topicId + " which doesn't exist in the topic metadata"); } } }); @@ -102,18 +102,20 @@ private Map> membersPerTopic(final AssignmentSpec assignmentS } /** - *

    The algorithm includes the following steps: + * The algorithm includes the following steps: *

      - *
    1. Generate a map of members per topic using the given member subscriptions.
      2. - *
    2. Generate a list of members called potentially unfilled members, which consists of members that have not met the minimum required quota of partitions for the assignment AND - * get a list called assigned sticky partitions for topic, which has the partitions that will be retained in the new assignment.
      3. - *
    3. Generate a list of unassigned partitions by calculating the difference between the total partitions for the topic and the assigned (sticky) partitions.
      4. + *
    4. Generate a map of members per topic using the given member subscriptions.
      5. + *
    5. Generate a list of members called potentially unfilled members, which consists of members that have not + * met the minimum required quota of partitions for the assignment AND get a list called assigned sticky + * partitions for topic, which has the partitions that will be retained in the new assignment.
      6. + *
    6. Generate a list of unassigned partitions by calculating the difference between the total partitions + * for the topic and the assigned (sticky) partitions.
      7. *
    7. Find members from the potentially unfilled members list that haven't met the total required quota - * i.e. minRequiredQuota + 1, if the member is designated to receive one of the excess partitions OR minRequiredQuota otherwise.
      8. + * i.e. minRequiredQuota + 1, if the member is designated to receive one of the excess partitions OR + * minRequiredQuota otherwise. *
    8. Assign partitions to them in ranges from the unassigned partitions per topic * based on the remaining partitions value.
      9. *
    - *

    */ @Override public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws PartitionAssignorException { @@ -195,7 +197,7 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit int unassignedPartitionsListStartPointer = 0; for (MemberWithRemainingAssignments pair : potentiallyUnfilledMembers) { String memberId = pair.memberId; - Integer remaining = pair.remaining; + int remaining = pair.remaining; if (numMembersWithExtraPartition > 0) { remaining++; numMembersWithExtraPartition--; @@ -211,6 +213,7 @@ public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws Partit } } }); + return new GroupAssignment(newAssignment); } } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 8aff7e9190e1c..91f6385f104e4 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -32,7 +32,7 @@ import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.assertThrows; public class RangeAssignorTest { private final RangeAssignor assignor = new RangeAssignor(); @@ -76,9 +76,9 @@ public void testOneConsumerSubscribedToNonExistentTopic() { ); AssignmentSpec assignmentSpec = new AssignmentSpec(members, topics); - GroupAssignment groupAssignment = assignor.assign(assignmentSpec); - assertTrue(groupAssignment.members().isEmpty()); + assertThrows(PartitionAssignorException.class, + () -> assignor.assign(assignmentSpec)); } @Test From 73dead3ee17c0179c5d23a0b1459e5af1c15c360 Mon Sep 17 00:00:00 2001 From: David Jacot Date: Wed, 10 May 2023 09:26:57 +0200 Subject: [PATCH 44/44] fix compilation --- .../apache/kafka/coordinator/group/RecordHelpersTest.java | 6 +++--- .../group/consumer/CurrentAssignmentBuilderTest.java | 4 ++-- .../group/consumer/TargetAssignmentBuilderTest.java | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java index d901309bfbd52..6504b57b2e465 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java @@ -44,9 +44,9 @@ import java.util.Map; import java.util.Set; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkSortedAssignment; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkSortedTopicAssignment; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkTopicAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkSortedAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkSortedTopicAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment; import static org.apache.kafka.coordinator.group.RecordHelpers.newCurrentAssignmentRecord; import static org.apache.kafka.coordinator.group.RecordHelpers.newCurrentAssignmentTombstoneRecord; import static org.apache.kafka.coordinator.group.RecordHelpers.newGroupEpochRecord; diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilderTest.java index b67c68e642534..bbe5cc5e0968e 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilderTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilderTest.java @@ -30,8 +30,8 @@ import java.util.Set; import java.util.stream.Stream; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkAssignment; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkTopicAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment; import static org.junit.jupiter.api.Assertions.assertEquals; public class CurrentAssignmentBuilderTest { diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java index a498af6ab8a40..7733f67b128eb 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java @@ -34,8 +34,8 @@ import java.util.Optional; import java.util.Set; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkAssignment; -import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkTopicAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment; +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment; import static org.apache.kafka.coordinator.group.RecordHelpers.newTargetAssignmentEpochRecord; import static org.apache.kafka.coordinator.group.RecordHelpers.newTargetAssignmentRecord; import static org.apache.kafka.coordinator.group.consumer.TargetAssignmentBuilder.createAssignmentMemberSpec;