dajac commented on code in PR #23723: URL: https://github.com/apache/kafka/pull/23723#discussion_r4398464203
########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignmentBuilderTest.java: ########## @@ -0,0 +1,402 @@ +/* + * 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.uniform2; + +import org.apache.kafka.common.Uuid; +import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.assignor.RangeSet; + +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.assertAssignment; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertSame; + +/** + * Tests of how {@link AssignmentBuilder} builds the result: which maps and sets of the current + * assignment it reuses, and which new sets it creates. A topic with {@code P} partitions and + * {@code N} subscribers has a minimum quota of {@code P / N} and {@code P % N} extra partitions. + */ +public class AssignmentBuilderTest { + private static final Uuid T1 = new Uuid(1L, 1L); + private static final Uuid T2 = new Uuid(2L, 1L); + private static final Uuid MISSING = new Uuid(9L, 1L); + + /** + * T1 has 6 partitions and 3 subscribers, so a minimum quota of 2 and no extra partition. A owns + * partitions 0 and 1, exactly its quota, and keeps its map. B owns partitions 2 to 5 and keeps + * 2 and 3, while C gets 4 and 5. Assigned again, the result keeps every map. + */ + @Test + public void testUnchangedMemberKeepsItsMap() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1), Map.of(T1, Set.of(0, 1))) + .withMember("B", Set.of(T1), Map.of(T1, Set.of(2, 3, 4, 5))) + .withMember("C", Set.of(T1)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 6) + .build(); + + var result = new AssignmentBuilder(spec, describer).build(); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1)), + "B", Map.of(T1, Set.of(2, 3)), + "C", Map.of(T1, Set.of(4, 5)) + ), result); + assertSame(spec.memberAssignment("A").partitions(), result.members().get("A").partitions()); + + var fedBack = GroupSpecFixture.after(spec, result).build(); + var again = new AssignmentBuilder(fedBack, describer).build(); + for (String id : fedBack.memberIds()) { + assertSame(fedBack.memberAssignment(id).partitions(), again.members().get(id).partitions(), id); + } + } + + /** + * A and C subscribe to T1, and B to T1 and T2. T1 has 6 partitions over its 3 subscribers, so a + * minimum quota of 2. T2 has 2 partitions and the single subscriber B, which owns both. B gives + * up partitions 4 and 5 of T1 to C, so it gets a new map. T2 does not change, so the new map + * has the same set of T2 that B holds. + */ + @Test + public void testChangedMemberSharesTheSetsOfItsUnchangedTopics() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1), Map.of(T1, Set.of(0, 1))) + .withMember("B", Set.of(T1, T2), Map.of(T1, Set.of(2, 3, 4, 5), T2, Set.of(0, 1))) + .withMember("C", Set.of(T1)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 6) + .withTopic(T2, 2) + .build(); + + var result = new AssignmentBuilder(spec, describer).build(); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1)), + "B", Map.of(T1, Set.of(2, 3), T2, Set.of(0, 1)), + "C", Map.of(T1, Set.of(4, 5)) + ), result); + assertSame(spec.memberAssignment("A").partitions(), result.members().get("A").partitions()); + assertNotSame(spec.memberAssignment("B").partitions(), result.members().get("B").partitions()); + assertSame(spec.memberAssignment("B").partitions().get(T2), result.members().get("B").partitions().get(T2)); + } + + /** + * A member that changes on one topic keeps its set of a later topic that it does not change + * on, even though partitions of that topic move. T1 has 6 partitions and T2 has 3, both over A, Review Comment: Not quite: there T2 did not move, so the builder did not look at it. Here T2 moves, and the builder has to see that A does not change on it. I merged both into one test with the two kinds of topics. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignmentBuilder.java: ########## @@ -0,0 +1,305 @@ +/* + * 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.uniform2; + +import org.apache.kafka.common.Uuid; +import org.apache.kafka.coordinator.group.Utils; +import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.api.assignor.GroupSpec; +import org.apache.kafka.coordinator.group.api.assignor.MemberAssignment; +import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber; +import org.apache.kafka.coordinator.group.assignor.RangeSet; +import org.apache.kafka.coordinator.group.assignor.uniform2.steps.AssignRemainingExtraPartitions; +import org.apache.kafka.coordinator.group.assignor.uniform2.steps.BalanceExtraPartitions; +import org.apache.kafka.coordinator.group.assignor.uniform2.steps.KeepStickyExtraPartitions; +import org.apache.kafka.coordinator.group.assignor.uniform2.util.IntList; +import org.apache.kafka.coordinator.group.modern.MemberAssignmentImpl; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Computes the assignment of a group. It reads the group and its current assignment, then runs + * {@link KeepStickyExtraPartitions}, {@link AssignRemainingExtraPartitions} and + * {@link BalanceExtraPartitions}, which decide the quotas. {@link PartitionPlacement} then decides + * which member gets every partition of every topic, and this class builds the assignments of the + * members from its decisions. + * + * <p>A member whose partitions do not change gets back the map it had, so that the coordinator + * finds it equal to its current assignment without comparing the sets. A member whose partitions + * change gets a new map, which shares the sets of its topics that do not change. A new set is a + * {@link RangeSet} when its partitions are consecutive, and a {@link HashSet} otherwise. + */ +public final class AssignmentBuilder { + /** + * The group spec. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * The members and topics of the group. + */ + private IndexedGroup group; + + /** + * The current assignment of the group. + */ + private CurrentAssignment currentAssignment; + + /** + * The quotas of the members, which the steps compute. + */ + private Quotas quotas; + + /** + * Per member, its new map once one of its topics changed, or null. + */ + private List<Map<Uuid, Set<Integer>>> newAssignments; + + /** + * Per member, the number of partitions of the topic being placed that it gets. All 0 between + * topics. + */ + private int[] placed; + + /** + * Per member, the lowest partition of the topic being placed that it gets. + */ + private int[] lowest; + + /** + * Per member, the highest partition of the topic being placed that it gets. + */ + private int[] highest; + + /** + * Per member, whether its partitions of the topic being placed change, because it gets a + * partition it does not own or loses one it owns. All false between topics. + */ + private boolean[] changed; + + /** + * The members whose partitions of the topic being placed change, to write their new sets and + * reset {@link #changed}. + */ + private IntList changedMembers; + + /** + * Per partition of the topic being placed, the next partition that its member gets. It links + * the partitions of every member from its lowest to its highest. + */ + private int[] nextPartition; + + /** + * The members getting partitions of the topic being placed, to reset {@link #placed} after + * the topic. + */ + private IntList membersGettingPartitions; + + /** + * @param groupSpec The group spec. + * @param subscribedTopicDescriber The describer of the subscribed topics. + */ + public AssignmentBuilder(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the group and its current assignment, computes the quotas, places the partitions of + * every topic, updates the assignments of the members whose partitions change, and builds the + * assignment of every member, reusing the unchanged maps. + * + * @return The assignment. + */ + public GroupAssignment build() { + // Reads the group and its current assignment, and computes the quotas. + group = new IndexedGroup.Builder(groupSpec, subscribedTopicDescriber).build(); + currentAssignment = new CurrentAssignment(groupSpec, group); + quotas = computeQuotas(); + + // Creates the new maps of the members, all null until one of their topics changes, and the + // state of updateMemberAssignments. + int memberCount = group.memberCount(); + newAssignments = new ArrayList<>(Collections.nCopies(memberCount, null)); + placed = new int[memberCount]; + lowest = new int[memberCount]; + highest = new int[memberCount]; + changed = new boolean[memberCount]; + changedMembers = new IntList(16); + membersGettingPartitions = new IntList(16); + nextPartition = new int[group.maxPartitionCount()]; + + // Places the partitions of every topic, and updates the assignments of the members whose + // partitions of the topic change. The assignment of the topic gives the member getting + // every partition. It is sized for the largest topic and reused for every topic. + var placement = new PartitionPlacement(group, currentAssignment, quotas); + var assignment = new int[group.maxPartitionCount()]; + for (int topic = 0; topic < group.topicCount(); topic++) { + if (placement.place(topic, assignment)) { + updateMemberAssignments(topic, assignment); + } + } + + // Builds the assignment of every member: its new map if one of its topics changed, and + // otherwise the map it holds, without its stale partitions if it has any. + Map<String, MemberAssignment> members = Utils.newHashMap(memberCount); + for (int member = 0; member < memberCount; member++) { + var partitions = newAssignments.get(member); + if (partitions == null) { + partitions = currentAssignment.holdsStalePartitions(member) + ? ownedPartitions(member) + : currentAssignment.assignment(member); + } Review Comment: Yes. A holds only stale partitions there and gets none of T1, so it gets a copy of its map without them, which is empty. `testStalePartitionsAreDroppedFromAMemberWhoseTopicsDoNotChange` goes through it too, with a member that keeps its partitions. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignmentBuilder.java: ########## @@ -0,0 +1,305 @@ +/* + * 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.uniform2; + +import org.apache.kafka.common.Uuid; +import org.apache.kafka.coordinator.group.Utils; +import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.api.assignor.GroupSpec; +import org.apache.kafka.coordinator.group.api.assignor.MemberAssignment; +import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber; +import org.apache.kafka.coordinator.group.assignor.RangeSet; +import org.apache.kafka.coordinator.group.assignor.uniform2.steps.AssignRemainingExtraPartitions; +import org.apache.kafka.coordinator.group.assignor.uniform2.steps.BalanceExtraPartitions; +import org.apache.kafka.coordinator.group.assignor.uniform2.steps.KeepStickyExtraPartitions; +import org.apache.kafka.coordinator.group.assignor.uniform2.util.IntList; +import org.apache.kafka.coordinator.group.modern.MemberAssignmentImpl; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Computes the assignment of a group. It reads the group and its current assignment, then runs + * {@link KeepStickyExtraPartitions}, {@link AssignRemainingExtraPartitions} and + * {@link BalanceExtraPartitions}, which decide the quotas. {@link PartitionPlacement} then decides + * which member gets every partition of every topic, and this class builds the assignments of the + * members from its decisions. + * + * <p>A member whose partitions do not change gets back the map it had, so that the coordinator + * finds it equal to its current assignment without comparing the sets. A member whose partitions + * change gets a new map, which shares the sets of its topics that do not change. A new set is a + * {@link RangeSet} when its partitions are consecutive, and a {@link HashSet} otherwise. + */ +public final class AssignmentBuilder { + /** + * The group spec. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * The members and topics of the group. + */ + private IndexedGroup group; + + /** + * The current assignment of the group. + */ + private CurrentAssignment currentAssignment; + + /** + * The quotas of the members, which the steps compute. + */ + private Quotas quotas; + + /** + * Per member, its new map once one of its topics changed, or null. + */ + private List<Map<Uuid, Set<Integer>>> newAssignments; + + /** + * Per member, the number of partitions of the topic being placed that it gets. All 0 between + * topics. + */ + private int[] placed; + + /** + * Per member, the lowest partition of the topic being placed that it gets. + */ + private int[] lowest; + + /** + * Per member, the highest partition of the topic being placed that it gets. + */ + private int[] highest; + + /** + * Per member, whether its partitions of the topic being placed change, because it gets a + * partition it does not own or loses one it owns. All false between topics. + */ + private boolean[] changed; + + /** + * The members whose partitions of the topic being placed change, to write their new sets and + * reset {@link #changed}. + */ + private IntList changedMembers; + + /** + * Per partition of the topic being placed, the next partition that its member gets. It links + * the partitions of every member from its lowest to its highest. + */ + private int[] nextPartition; + + /** + * The members getting partitions of the topic being placed, to reset {@link #placed} after + * the topic. + */ + private IntList membersGettingPartitions; + + /** + * @param groupSpec The group spec. + * @param subscribedTopicDescriber The describer of the subscribed topics. + */ + public AssignmentBuilder(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the group and its current assignment, computes the quotas, places the partitions of + * every topic, updates the assignments of the members whose partitions change, and builds the + * assignment of every member, reusing the unchanged maps. + * + * @return The assignment. + */ + public GroupAssignment build() { + // Reads the group and its current assignment, and computes the quotas. + group = new IndexedGroup.Builder(groupSpec, subscribedTopicDescriber).build(); + currentAssignment = new CurrentAssignment(groupSpec, group); + quotas = computeQuotas(); + + // Creates the new maps of the members, all null until one of their topics changes, and the + // state of updateMemberAssignments. + int memberCount = group.memberCount(); + newAssignments = new ArrayList<>(Collections.nCopies(memberCount, null)); + placed = new int[memberCount]; + lowest = new int[memberCount]; + highest = new int[memberCount]; + changed = new boolean[memberCount]; + changedMembers = new IntList(16); + membersGettingPartitions = new IntList(16); + nextPartition = new int[group.maxPartitionCount()]; + + // Places the partitions of every topic, and updates the assignments of the members whose + // partitions of the topic change. The assignment of the topic gives the member getting + // every partition. It is sized for the largest topic and reused for every topic. + var placement = new PartitionPlacement(group, currentAssignment, quotas); + var assignment = new int[group.maxPartitionCount()]; + for (int topic = 0; topic < group.topicCount(); topic++) { + if (placement.place(topic, assignment)) { + updateMemberAssignments(topic, assignment); + } + } + + // Builds the assignment of every member: its new map if one of its topics changed, and + // otherwise the map it holds, without its stale partitions if it has any. + Map<String, MemberAssignment> members = Utils.newHashMap(memberCount); + for (int member = 0; member < memberCount; member++) { + var partitions = newAssignments.get(member); + if (partitions == null) { + partitions = currentAssignment.holdsStalePartitions(member) + ? ownedPartitions(member) + : currentAssignment.assignment(member); + } + members.put(group.memberId(member), new MemberAssignmentImpl(partitions)); + } + return new GroupAssignment(members); + } + + /** + * @return The quotas, computed by running the steps in order. + */ + private Quotas computeQuotas() { + var quotas = new Quotas(group); + new KeepStickyExtraPartitions(group, currentAssignment, quotas).run(); + new AssignRemainingExtraPartitions(group, quotas).run(); + new BalanceExtraPartitions(group, currentAssignment, quotas).run(); + return quotas; + } + + /** + * Updates the new assignments of the members whose partitions of the topic change. Such a + * member gets a new set, or loses the topic if it gets none of the topic's partitions. Every + * other member gets exactly the partitions it owns, so it keeps its set. + */ + private void updateMemberAssignments(int topic, int[] assignment) { + int partitionCount = group.partitionCount(topic); + var owners = currentAssignment.owners(topic); + membersGettingPartitions.clear(); + + // Counts the partitions of every member and links them, from its lowest to its highest. A + // partition going to a member which does not own it changes both that member and its + // owner, if it has one. + for (int partition = 0; partition < partitionCount; partition++) { + int member = assignment[partition]; + if (placed[member]++ == 0) { + membersGettingPartitions.add(member); + lowest[member] = partition; + } else { + nextPartition[highest[member]] = partition; + } + highest[member] = partition; + int owner = owners == null ? -1 : owners[partition]; + if (owner != member) { + markChanged(member); + if (owner >= 0) { + markChanged(owner); + } + } + } + + // Gives every changed member its new set, or removes the topic from it if it gets none of + // the topic's partitions. + var topicId = group.topicId(topic); + for (int i = 0; i < changedMembers.size(); i++) { + int member = changedMembers.get(i); + if (placed[member] == 0) { + newAssignment(member).remove(topicId); + } else { + newAssignment(member).put(topicId, newSet(member)); + } + changed[member] = false; + } + changedMembers.clear(); + + // Resets the counts of the members getting partitions for the next topic, at most one per + // partition of the topic rather than every member of the group. + for (int i = 0; i < membersGettingPartitions.size(); i++) { + placed[membersGettingPartitions.get(i)] = 0; + } + } + + /** + * Marks the partitions of the topic of the member as changed. The member is added to + * {@link #changedMembers} the first time. + */ + private void markChanged(int member) { + if (!changed[member]) { + changed[member] = true; + changedMembers.add(member); + } + } + + /** + * @return The new set of the partitions of the topic being placed that the member gets: a + * {@link RangeSet} when they are consecutive, and otherwise a hash set filled by + * following {@link #nextPartition} from the lowest partition of the member to its + * highest. + */ + private Set<Integer> newSet(int member) { + if (highest[member] - lowest[member] + 1 == placed[member]) { + return new RangeSet(lowest[member], highest[member] + 1); + } + Set<Integer> partitions = Utils.newHashSet(placed[member]); + int partition = lowest[member]; + partitions.add(partition); + while (partition != highest[member]) { + partition = nextPartition[partition]; + partitions.add(partition); + } + return partitions; + } Review Comment: Thanks! Agreed. Replaying the records rebuilds hash sets, so we lose it there. We could consider keeping ranges on replay if possible. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/Quotas.java: ########## @@ -131,6 +138,15 @@ public int assignmentSize(int member) { return minAssignmentSizes[group.cohortOf(member)] + extraPartitionCounts[member]; } + /** + * @return An upper bound of the number of topics of which the member gets partitions. A topic + * with a minimum quota above 0 counts twice when the member gets one of its extra + * partitions. + */ + public int maxTopicsWithPartitions(int member) { + return topicsWithMinQuotaAboveZero[group.cohortOf(member)] + extraPartitionCounts[member]; + } Review Comment: Good point. It only sized the new map of a member whose partitions change, so that it does not grow, and that map rarely grows by more than a few entries. I removed it. ########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignmentBuilderTest.java: ########## @@ -0,0 +1,402 @@ +/* + * 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.uniform2; + +import org.apache.kafka.common.Uuid; +import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.assignor.RangeSet; + +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.assertAssignment; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertSame; + +/** + * Tests of how {@link AssignmentBuilder} builds the result: which maps and sets of the current + * assignment it reuses, and which new sets it creates. A topic with {@code P} partitions and + * {@code N} subscribers has a minimum quota of {@code P / N} and {@code P % N} extra partitions. + */ +public class AssignmentBuilderTest { + private static final Uuid T1 = new Uuid(1L, 1L); + private static final Uuid T2 = new Uuid(2L, 1L); + private static final Uuid MISSING = new Uuid(9L, 1L); + + /** + * T1 has 6 partitions and 3 subscribers, so a minimum quota of 2 and no extra partition. A owns + * partitions 0 and 1, exactly its quota, and keeps its map. B owns partitions 2 to 5 and keeps + * 2 and 3, while C gets 4 and 5. Assigned again, the result keeps every map. + */ + @Test + public void testUnchangedMemberKeepsItsMap() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1), Map.of(T1, Set.of(0, 1))) + .withMember("B", Set.of(T1), Map.of(T1, Set.of(2, 3, 4, 5))) + .withMember("C", Set.of(T1)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 6) + .build(); + + var result = new AssignmentBuilder(spec, describer).build(); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1)), + "B", Map.of(T1, Set.of(2, 3)), + "C", Map.of(T1, Set.of(4, 5)) + ), result); + assertSame(spec.memberAssignment("A").partitions(), result.members().get("A").partitions()); + + var fedBack = GroupSpecFixture.after(spec, result).build(); + var again = new AssignmentBuilder(fedBack, describer).build(); + for (String id : fedBack.memberIds()) { + assertSame(fedBack.memberAssignment(id).partitions(), again.members().get(id).partitions(), id); + } + } + + /** + * A and C subscribe to T1, and B to T1 and T2. T1 has 6 partitions over its 3 subscribers, so a + * minimum quota of 2. T2 has 2 partitions and the single subscriber B, which owns both. B gives + * up partitions 4 and 5 of T1 to C, so it gets a new map. T2 does not change, so the new map + * has the same set of T2 that B holds. + */ + @Test + public void testChangedMemberSharesTheSetsOfItsUnchangedTopics() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1), Map.of(T1, Set.of(0, 1))) + .withMember("B", Set.of(T1, T2), Map.of(T1, Set.of(2, 3, 4, 5), T2, Set.of(0, 1))) + .withMember("C", Set.of(T1)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 6) + .withTopic(T2, 2) + .build(); + + var result = new AssignmentBuilder(spec, describer).build(); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1)), + "B", Map.of(T1, Set.of(2, 3), T2, Set.of(0, 1)), + "C", Map.of(T1, Set.of(4, 5)) + ), result); + assertSame(spec.memberAssignment("A").partitions(), result.members().get("A").partitions()); + assertNotSame(spec.memberAssignment("B").partitions(), result.members().get("B").partitions()); + assertSame(spec.memberAssignment("B").partitions().get(T2), result.members().get("B").partitions().get(T2)); + } + + /** + * A member that changes on one topic keeps its set of a later topic that it does not change + * on, even though partitions of that topic move. T1 has 6 partitions and T2 has 3, both over A, + * B and C, so minimum quotas of 2 and 1 and no extra partition. On T1, A keeps partitions 0 and + * 1 of the 4 it owns, and C gets 2 and 3. On T2, A keeps partition 0, B keeps 1, and C gets 2, + * which B gives up. A changes on T1 only, so its new map has the same set of T2 that A holds, + * and B changes on T2 only, so its new map has the same set of T1 that B holds. + */ + @Test + public void testChangedMemberSharesItsSetOfAMovingTopicItDoesNotChangeOn() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2), Map.of(T1, Set.of(0, 1, 2, 3), T2, Set.of(0))) + .withMember("B", Set.of(T1, T2), Map.of(T1, Set.of(4, 5), T2, Set.of(1, 2))) + .withMember("C", Set.of(T1, T2)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 6) + .withTopic(T2, 3) + .build(); + + var result = new AssignmentBuilder(spec, describer).build(); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1), T2, Set.of(0)), + "B", Map.of(T1, Set.of(4, 5), T2, Set.of(1)), + "C", Map.of(T1, Set.of(2, 3), T2, Set.of(2)) + ), result); + assertSame(spec.memberAssignment("A").partitions().get(T2), result.members().get("A").partitions().get(T2)); + assertSame(spec.memberAssignment("B").partitions().get(T1), result.members().get("B").partitions().get(T1)); + } + + /** + * Nobody owns the 4 partitions of T1, so its 2 subscribers both change: A gets partitions 0 and + * 1, and B gets 2 and 3, each in a new map with a {@link RangeSet}. + */ + @Test + public void testNewSetsOfATopicWithoutOwnersAreRangeSets() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1)) + .withMember("B", Set.of(T1)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 4) + .build(); + + var result = new AssignmentBuilder(spec, describer).build(); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1)), + "B", Map.of(T1, Set.of(2, 3)) + ), result); + assertEquals(Map.of( + "A", Map.of(T1, RangeSet.class), + "B", Map.of(T1, RangeSet.class) + ), setClasses(result)); + } + + /** + * T1 has 6 partitions and 3 subscribers, so a minimum quota of 2. A owns partitions 0 to 2 and + * keeps 0 and 1, B owns 3 to 5 and keeps 3 and 4, and C gets partitions 2 and 5. The new sets + * of A and B are consecutive, and that of C is not. + */ + @Test + public void testNewSetsAreRangeSetsWhenConsecutive() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1), Map.of(T1, Set.of(0, 1, 2))) + .withMember("B", Set.of(T1), Map.of(T1, Set.of(3, 4, 5))) + .withMember("C", Set.of(T1)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 6) + .build(); + + var result = new AssignmentBuilder(spec, describer).build(); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1)), + "B", Map.of(T1, Set.of(3, 4)), + "C", Map.of(T1, Set.of(2, 5)) + ), result); + assertEquals(Map.of( + "A", Map.of(T1, RangeSet.class), + "B", Map.of(T1, RangeSet.class), + "C", Map.of(T1, HashSet.class) + ), setClasses(result)); + } + + /** + * T1 has 6 partitions and 3 subscribers, so a minimum quota of 2. A owns partitions 0, 2, 3 and + * 5, and keeps 0 and 2, which are not consecutive. B gets 1 and 3, which are not consecutive + * either, and C gets 4 and 5. + */ + @Test + public void testNewSetsAreHashSetsWhenNotConsecutive() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1), Map.of(T1, Set.of(0, 2, 3, 5))) + .withMember("B", Set.of(T1)) + .withMember("C", Set.of(T1)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 6) + .build(); + + var result = new AssignmentBuilder(spec, describer).build(); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 2)), + "B", Map.of(T1, Set.of(1, 3)), + "C", Map.of(T1, Set.of(4, 5)) + ), result); + assertEquals(Map.of( + "A", Map.of(T1, HashSet.class), + "B", Map.of(T1, HashSet.class), + "C", Map.of(T1, RangeSet.class) + ), setClasses(result)); + } + + /** + * T1 has 5 partitions and 2 subscribers, so a minimum quota of 2 and 1 extra partition, A's. A + * owns partitions 0 and 1, keeps them and gets 2, and B gets 3 and 4. Both new sets are + * consecutive. + */ + @Test + public void testNewSetOfAMemberKeepingAndGettingPartitionsIsARangeSetWhenConsecutive() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1), Map.of(T1, Set.of(0, 1))) + .withMember("B", Set.of(T1)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 5) + .build(); + + var result = new AssignmentBuilder(spec, describer).build(); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1, 2)), + "B", Map.of(T1, Set.of(3, 4)) + ), result); + assertEquals(Map.of( + "A", Map.of(T1, RangeSet.class), + "B", Map.of(T1, RangeSet.class) + ), setClasses(result)); + } + + /** + * T1 has 6 partitions and 2 subscribers, so a minimum quota of 3. B owns partitions 1, 3 and 5, + * its quota, and keeps its map. A owns 0 and 4, keeps them and gets 2, which nobody owns, so + * its new set is 0, 2 and 4. + */ + @Test + public void testNewSetOfAMemberKeepingAndGettingPartitionsHasThemAll() { Review Comment: It checked that a set that is not consecutive has the partitions that the member keeps as well as those it gets, which `testNewSetsAreHashSetsWhenNotConsecutive` already covers. I removed it. ########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/Uniform2AssignorTest.java: ########## @@ -0,0 +1,758 @@ +/* + * 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.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage; +import org.apache.kafka.coordinator.common.runtime.MetadataImageBuilder; +import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType; +import org.apache.kafka.coordinator.group.assignor.uniform2.GroupSpecFixture; +import org.apache.kafka.coordinator.group.assignor.uniform2.TopicsFixture; +import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl; +import org.apache.kafka.coordinator.group.modern.TopicIds; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; + +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.apache.kafka.coordinator.group.AssignmentTestUtil.assertAssignment; +import static org.apache.kafka.coordinator.group.assignor.uniform2.AssignmentTestUtils.assertValidAssignment; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; + +/** + * End to end tests of the uniform2 assignor, with expectations traced by hand. + */ +public class Uniform2AssignorTest { + private static final Uuid T1 = new Uuid(1L, 1L); + private static final Uuid T2 = new Uuid(2L, 1L); + private static final Uuid T3 = new Uuid(3L, 1L); + private static final Uuid T4 = new Uuid(4L, 1L); + private static final Uuid T5 = new Uuid(5L, 1L); + private static final Uuid MISSING = new Uuid(9L, 1L); + + private final Uniform2Assignor assignor = new Uniform2Assignor(); + + @Test + public void testName() { + assertEquals("uniform2", assignor.name()); + } + + @Test + public void testEmptyGroup() { + var result = assignor.assign(new GroupSpecFixture().build(), new TopicsFixture().withTopic(T1, 3).build()); + assertEquals(Map.of(), result.members()); + } + + @Test + public void testMissingAndUnsubscribedTopicsAreIgnored() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(MISSING)) + .withMember("B", Set.of(MISSING)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 3) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(), + "B", Map.of() + ), result); + } + + /** + * A subscribes to T1 and holds partitions 0 and 1 of T1, partition 0 of T2, which only B + * subscribes to, and partition 0 of a topic that does not exist. T1, of 4 partitions over A and + * B, has a minimum quota of 2, which both own, so T1 does not change. A loses its partitions of + * T2 and of the missing topic, and B gets the single partition of T2. + */ + @Test + public void testStalePartitionsAreDropped() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1), Map.of(T1, Set.of(0, 1), T2, Set.of(0), MISSING, Set.of(0))) + .withMember("B", Set.of(T1, T2), Map.of(T1, Set.of(2, 3))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 4) + .withTopic(T2, 1) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1)), + "B", Map.of(T1, Set.of(2, 3), T2, Set.of(0)) + ), result); + } + + /** + * T1 has 2 partitions, so a minimum quota of 0 and 2 extra partitions, T2 has 5, so a minimum + * quota of 1 and 2 extra partitions, and T3 has 7, so a minimum quota of 2 and 1 extra + * partition. The 5 extra partitions go round robin, T1's to A and B, T2's to C and A, and T3's + * to B. The members below their quota then take the partitions of every topic in partition + * order. + */ + @Test + public void testHomogeneousSpreadsEveryTopic() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2, T3)) + .withMember("B", Set.of(T1, T2, T3)) + .withMember("C", Set.of(T1, T2, T3)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 2) + .withTopic(T2, 5) + .withTopic(T3, 7) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0), T2, Set.of(0, 1), T3, Set.of(0, 1)), + "B", Map.of(T1, Set.of(1), T2, Set.of(2), T3, Set.of(2, 3, 4)), + "C", Map.of(T2, Set.of(3, 4), T3, Set.of(5, 6)) + ), result); + } + + /** + * A and D subscribe to T2, and B and C to T1 and T2, so that A and D form a cohort and B and C + * another, although B and C come between A and D in member order. T1, of 2 partitions over B + * and C, has a minimum quota of 1, and T2, of 8 partitions over the four members, a minimum + * quota of 2. Nobody owns anything, and the partitions of T2 go cohort by cohort: 0 to 3 to A + * and D, then 4 to 7 to B and C. + */ + @Test + public void testMembersWithTheSameSubscriptionShareACohort() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T2)) + .withMember("B", Set.of(T1, T2)) + .withMember("C", Set.of(T1, T2)) + .withMember("D", Set.of(T2)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 2) + .withTopic(T2, 8) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T2, Set.of(0, 1)), + "B", Map.of(T1, Set.of(0), T2, Set.of(4, 5)), + "C", Map.of(T1, Set.of(1), T2, Set.of(6, 7)), + "D", Map.of(T2, Set.of(2, 3)) + ), result); + } + + /** + * C leaves the assignment of {@link #testHomogeneousSpreadsEveryTopic}. T1 now has a + * minimum quota of 1, which A and B hold. T2 has a minimum quota of 2 and 1 extra partition, T3 + * has a minimum quota of 3 and 1 extra partition, and neither A nor B owns more than the + * minimum quota of them, so their extra partitions are assigned, T2's to A and T3's to B. A and + * B then take the partitions of C in partition order, and nothing moves between them. + */ + @Test + public void testLeaveOnlyMovesThePartitionsOfTheLeaver() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2, T3), Map.of(T1, Set.of(0), T2, Set.of(0, 1), T3, Set.of(0, 1))) + .withMember("B", Set.of(T1, T2, T3), Map.of(T1, Set.of(1), T2, Set.of(2), T3, Set.of(2, 3, 4))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 2) + .withTopic(T2, 5) + .withTopic(T3, 7) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0), T2, Set.of(0, 1, 3), T3, Set.of(0, 1, 5)), + "B", Map.of(T1, Set.of(1), T2, Set.of(2, 4), T3, Set.of(2, 3, 4, 6)) + ), result); + } + + /** + * A member joining a stable homogeneous group only takes its quota. With 3 topics of 6 + * partitions and 2 members, A has partitions 0 to 2 of each topic and B 3 to 5. With a third + * member, each member has 2 partitions of each topic, so A and B keep their 2 lowest partitions + * of each topic and give the third one to the joiner. Nothing else moves. + */ + @Test + public void testJoinOnlyMovesThePartitionsOfTheJoiner() { + var describer = new TopicsFixture() + .withTopic(T1, 6) + .withTopic(T2, 6) + .withTopic(T3, 6) + .build(); + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2, T3)) + .withMember("B", Set.of(T1, T2, T3)) + .build(); + var before = assignor.assign(spec, describer); + var joined = GroupSpecFixture.after(spec, before) + .withMember("C", Set.of(T1, T2, T3)) + .build(); + + var after = assignor.assign(joined, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1), T2, Set.of(0, 1), T3, Set.of(0, 1)), + "B", Map.of(T1, Set.of(3, 4), T2, Set.of(3, 4), T3, Set.of(3, 4)), + "C", Map.of(T1, Set.of(2, 5), T2, Set.of(2, 5), T3, Set.of(2, 5)) + ), after); + } + + /** + * A and B subscribe to T1 and T2. T1 had 4 partitions, of which A holds 0 and 1 and B holds 2 + * and 3, and now has 6, so a minimum quota of 3. A gets the new partition 4, and B gets 5. T2, + * of 2 partitions, does not change, and both keep their set of it. + */ + @Test + public void testAddedPartitionsGoToTheMembersBelowTheirQuota() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2), Map.of(T1, Set.of(0, 1), T2, Set.of(0))) + .withMember("B", Set.of(T1, T2), Map.of(T1, Set.of(2, 3), T2, Set.of(1))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 6) + .withTopic(T2, 2) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1, 4), T2, Set.of(0)), + "B", Map.of(T1, Set.of(2, 3, 5), T2, Set.of(1)) + ), result); + } + + /** + * A, B and C subscribed to T1, of 6 partitions, and hold 2 partitions each. C now subscribes + * to T2, of 3 partitions, instead. Its partitions 4 and 5 of T1 are dropped and go to A and B, + * whose minimum quota of T1 is now 3, and C gets all the partitions of T2. + */ + @Test + public void testMemberChangingItsSubscriptionGivesUpTheTopicItLeaves() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1), Map.of(T1, Set.of(0, 1))) + .withMember("B", Set.of(T1), Map.of(T1, Set.of(2, 3))) + .withMember("C", Set.of(T2), Map.of(T1, Set.of(4, 5))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 6) + .withTopic(T2, 3) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1, 4)), + "B", Map.of(T1, Set.of(2, 3, 5)), + "C", Map.of(T2, Set.of(0, 1, 2)) + ), result); + } + + /** + * A subscribes to T1, and B and C to T1 and T2. T1, of 5 partitions, has a minimum quota of 1 + * and 2 extra partitions. A owns partitions 0 and 1 and keeps one extra partition. A is the + * only member of its cohort and already has an extra partition of T1, so its cohort gets no + * other, and the second one goes to B, the first of the other cohort. B gets partitions 2 and 3 + * of T1, and C gets 4. T2, of 2 partitions over B and C, gives each of them 1. + */ + @Test + public void testCohortWhoseMembersAllHaveAnExtraPartitionGetsNoMore() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1), Map.of(T1, Set.of(0, 1))) + .withMember("B", Set.of(T1, T2)) + .withMember("C", Set.of(T1, T2)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 5) + .withTopic(T2, 2) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1)), + "B", Map.of(T1, Set.of(2, 3), T2, Set.of(0)), + "C", Map.of(T1, Set.of(4), T2, Set.of(1)) + ), result); + } + + /** + * A, B and C subscribe to T1 and T2, of 2 partitions each, and T3, of 5 partitions. T1 and T2 + * have a minimum quota of 0 and 2 extra partitions each, which B and C keep, as each owns one + * partition of each. T3 has a minimum quota of 1 and 2 extra partitions. A owns partitions 0 + * and 1 and keeps one. The second goes to B: A has the fewest extra partitions but already has + * one of T3, and B is the first of B and C, which have 2 each. B, then at 4 partitions against + * 2 for A, gives A its extra partition of T2, since A cannot take another one of T3. A gets + * partition 0 of T2, and B gets partition 4 of T3, which nobody owns. + */ + @Test + public void testAnExtraPartitionGoesToAnotherMemberWhenTheSmallestHaveOneOfTheTopic() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2, T3), Map.of(T3, Set.of(0, 1))) + .withMember("B", Set.of(T1, T2, T3), Map.of(T1, Set.of(0), T2, Set.of(0), T3, Set.of(2))) + .withMember("C", Set.of(T1, T2, T3), Map.of(T1, Set.of(1), T2, Set.of(1), T3, Set.of(3))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 2) + .withTopic(T2, 2) + .withTopic(T3, 5) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T2, Set.of(0), T3, Set.of(0, 1)), + "B", Map.of(T1, Set.of(0), T3, Set.of(2, 4)), + "C", Map.of(T1, Set.of(1), T2, Set.of(1), T3, Set.of(3)) + ), result); + } + + /** + * A subscribes to T1, B to T1 and T3, and C and D to T2 and T4. T1 and T2, of 1 partition each + * over 2 subscribers, have a minimum quota of 0 and 1 extra partition each. T3 gives B 1 + * partition, and T4 gives C and D 2 each. Nobody owns anything. The extra partition of T1 goes + * to A, whose assignment is smaller than that of B. The extra partition of T2 goes to C, the + * first of its subscribers. B, which is smaller than C and could have taken the extra + * partition of T1, does not subscribe to T2. + */ + @Test + public void testCohortsWaitingForAnExtraPartitionOfOneTopicDoNotCarryOverToTheNext() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1)) + .withMember("B", Set.of(T1, T3)) + .withMember("C", Set.of(T2, T4)) + .withMember("D", Set.of(T2, T4)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 1) + .withTopic(T2, 1) + .withTopic(T3, 1) + .withTopic(T4, 4) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0)), + "B", Map.of(T3, Set.of(0)), + "C", Map.of(T2, Set.of(0), T4, Set.of(0, 1)), + "D", Map.of(T4, Set.of(2, 3)) + ), result); + } + + /** + * A subscribes to T1 and T2, B to T2 only. T1, 4 partitions, has a single subscriber, so A has + * a minimum quota of 4. T2, 3 partitions, has a minimum quota of 1, and its extra partition + * goes to B, whose assignment is smaller. + */ + @Test + public void testHeterogeneousBalancesAcrossTopics() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2)) + .withMember("B", Set.of(T2)) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 4) + .withTopic(T2, 3) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1, 2, 3), T2, Set.of(0)), + "B", Map.of(T2, Set.of(1, 2)) + ), result); + } + + /** + * A subscribes to T1 and T2, B to T1 and C to T2, of 3 partitions each, so each topic has a + * minimum quota of 1 and 1 extra partition. A owns 2 partitions of both, and B and C one each, + * so A keeps both extra partitions and has 4 partitions against 1. The balance step then moves + * the extra partition of T1 to B and the one of T2 to C. A gives up one partition of each, the + * fewest revocations for assignment sizes within one. + */ + @Test + public void testBalanceMovesExtraPartitionsKeptByTheirOwner() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2), Map.of(T1, Set.of(0, 1), T2, Set.of(0, 1))) + .withMember("B", Set.of(T1), Map.of(T1, Set.of(2))) + .withMember("C", Set.of(T2), Map.of(T2, Set.of(2))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 3) + .withTopic(T2, 3) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0), T2, Set.of(0)), + "B", Map.of(T1, Set.of(1, 2)), + "C", Map.of(T2, Set.of(1, 2)) + ), result); + } + + /** + * A subscribes to T1, T2 and T3, B to T1 and T2, C to T2 and T3, and D to T1 and T3, in four + * cohorts. T1, T2 and T3, of 5 partitions each over 3 subscribers, have a minimum quota of 1 + * and 2 extra partitions each, all kept by owners of 2 partitions: A keeps one of every topic, + * B one of T1 and one of T2, and C one of T3. The assignment sizes are 6, 4, 3 and 2. A gives + * the extra partition of T3 to D, the smallest, then the one of T2 to C. The assignment sizes + * end at 4, 4, 4 and 3, and no extra partition can move to a subscriber of its topic at least + * two partitions smaller. C gets partition 1 of T2, and D partition 1 of T3. + */ + @Test + public void testBalanceMovesUntilNoExtraPartitionCanMove() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2, T3), Map.of(T1, Set.of(0, 1), T2, Set.of(0, 1), T3, Set.of(0, 1))) + .withMember("B", Set.of(T1, T2), Map.of(T1, Set.of(2, 3), T2, Set.of(2, 3))) + .withMember("C", Set.of(T2, T3), Map.of(T2, Set.of(4), T3, Set.of(2, 3))) + .withMember("D", Set.of(T1, T3), Map.of(T1, Set.of(4), T3, Set.of(4))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 5) + .withTopic(T2, 5) + .withTopic(T3, 5) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1), T2, Set.of(0), T3, Set.of(0)), + "B", Map.of(T1, Set.of(2, 3), T2, Set.of(2, 3)), + "C", Map.of(T2, Set.of(1, 4), T3, Set.of(2, 3)), + "D", Map.of(T1, Set.of(4), T3, Set.of(1, 4)) + ), result); + } + + /** + * A subscribes to T1 and T4, and B and C to T1, T2 and T3. T1, of 5 partitions over the three, + * has a minimum quota of 1 and 2 extra partitions, kept by A and B, which own 2 partitions + * each. T2 and T3, of 3 partitions over B and C, have 1 extra partition each, both kept by C. + * T4 gives A its 5 partitions. The assignment sizes are 7, 4 and 5. B, the smallest, already + * has an extra partition of T1, the only topic of the extra partitions of A, so C, the next, + * takes it. C, at 6, then gives B its extra partition of T3. A gives up partition 1 of T1 to + * C, and C gives up partition 2 of T3 to B. + */ + @Test + public void testTheSmallestMemberThatCannotTakeIsPassedOver() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T4), Map.of(T1, Set.of(0, 1), T4, Set.of(0, 1, 2, 3, 4))) + .withMember("B", Set.of(T1, T2, T3), Map.of(T1, Set.of(2, 3), T2, Set.of(0), T3, Set.of(0))) + .withMember("C", Set.of(T1, T2, T3), Map.of(T1, Set.of(4), T2, Set.of(1, 2), T3, Set.of(1, 2))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 5) + .withTopic(T2, 3) + .withTopic(T3, 3) + .withTopic(T4, 5) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0), T4, Set.of(0, 1, 2, 3, 4)), + "B", Map.of(T1, Set.of(2, 3), T2, Set.of(0), T3, Set.of(0, 2)), + "C", Map.of(T1, Set.of(1, 4), T2, Set.of(1, 2), T3, Set.of(1)) + ), result); + } + + /** + * A subscribes to T1 and T3, B to T1 and T2, and C to T2. T1, of 3 partitions over A and B, has + * 1 extra partition, which A keeps as it owns 2 partitions. T2, of 3 partitions over B and C, + * has 1, which B keeps for the same reason. T3 gives A its 2 partitions. The assignment sizes + * are 4, 3 and 1. A, the largest, has no receiver, since B is only one partition smaller, so A + * waits. B gives its extra partition of T2 to C, and is then two partitions smaller than A, + * which wakes up and gives B its extra partition of T1. B gets partition 1 of T1, and C + * partition 1 of T2. + */ + @Test + public void testAWaitingGiverWakesUpWhenASubscriberOfItsTopicsGetsSmaller() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T3), Map.of(T1, Set.of(0, 1), T3, Set.of(0, 1))) + .withMember("B", Set.of(T1, T2), Map.of(T1, Set.of(2), T2, Set.of(0, 1))) + .withMember("C", Set.of(T2), Map.of(T2, Set.of(2))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 3) + .withTopic(T2, 3) + .withTopic(T3, 2) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0), T3, Set.of(0, 1)), + "B", Map.of(T1, Set.of(1, 2), T2, Set.of(0)), + "C", Map.of(T2, Set.of(1, 2)) + ), result); + } + + /** + * A and C subscribe to T1 to T4, and B to T1 to T5. T1 to T4, of 4 partitions over the three, + * have a minimum quota of 1 and 1 extra partition each, and T5, of 1 partition, gives B 1 more. + * C owns 2 partitions of T1 and keeps its extra partition. A owns 2 partitions of T2, T3 and T4 + * and keeps their extra partitions. C also owns 2 partitions of T3, but A keeps the extra + * partition of T3, since A and C are then equally large and A comes first. The assignment + * sizes are 7, 5 and 5. B and C are equally small, and C receives first, since it owns more + * than the minimum quota of T3. A gives C the extra partition of T3, which revokes no partition + * from C, rather than that of T4, its last topic. B gets partition 1 of T3, which A gives up. + */ + @Test + public void testTheGiverGivesTheExtraPartitionOfATopicTheReceiverOwnsMoreThanTheMinimumQuotaOf() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2, T3, T4), Map.of(T1, Set.of(2), T2, Set.of(0, 1), T3, Set.of(0, 1), T4, Set.of(0, 1))) + .withMember("B", Set.of(T1, T2, T3, T4, T5), Map.of(T1, Set.of(3), T2, Set.of(2), T4, Set.of(2), T5, Set.of(0))) + .withMember("C", Set.of(T1, T2, T3, T4), Map.of(T1, Set.of(0, 1), T2, Set.of(3), T3, Set.of(2, 3), T4, Set.of(3))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 4) + .withTopic(T2, 4) + .withTopic(T3, 4) + .withTopic(T4, 4) + .withTopic(T5, 1) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(2), T2, Set.of(0, 1), T3, Set.of(0), T4, Set.of(0, 1)), + "B", Map.of(T1, Set.of(3), T2, Set.of(2), T3, Set.of(1), T4, Set.of(2), T5, Set.of(0)), + "C", Map.of(T1, Set.of(0, 1), T2, Set.of(3), T3, Set.of(2, 3), T4, Set.of(3)) + ), result); + } + + /** + * A and B subscribe to T1 to T4, of 3 partitions each, so a minimum quota of 1 and 1 extra + * partition each. B owns 2 partitions of every topic and keeps the four extra partitions, so + * the assignment sizes are 4 and 8. B gives A two extra partitions in a row, of its last topics + * first: T4, then T3. A gets partition 1 of T3 and of T4, which B gives up. + */ + @Test + public void testAGiverGivesSeveralExtraPartitionsInARow() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2, T3, T4), Map.of(T1, Set.of(2), T2, Set.of(2), T3, Set.of(2), T4, Set.of(2))) + .withMember("B", Set.of(T1, T2, T3, T4), Map.of(T1, Set.of(0, 1), T2, Set.of(0, 1), T3, Set.of(0, 1), T4, Set.of(0, 1))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 3) + .withTopic(T2, 3) + .withTopic(T3, 3) + .withTopic(T4, 3) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(2), T2, Set.of(2), T3, Set.of(1, 2), T4, Set.of(1, 2)), + "B", Map.of(T1, Set.of(0, 1), T2, Set.of(0, 1), T3, Set.of(0), T4, Set.of(0)) + ), result); + } + + /** + * An assignment that already has the properties is kept as it is, even though the assignor + * would not have computed it from scratch. T1 has a minimum quota of 2. T2 has a minimum quota + * of 1 and 1 extra partition, which A keeps since it owns 2 partitions of T2. The assignment + * sizes are 4 and 3. From scratch, A would get partitions 0 and 1 of T1 rather than 2 and 3. + */ + @Test + public void testAssignmentWithThePropertiesDoesNotChange() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2), Map.of(T1, Set.of(2, 3), T2, Set.of(1, 2))) + .withMember("B", Set.of(T1, T2), Map.of(T1, Set.of(0, 1), T2, Set.of(0))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 4) + .withTopic(T2, 3) + .build(); + + var result = assignor.assign(spec, describer); + + for (String id : spec.memberIds()) { + assertSame(spec.memberAssignment(id).partitions(), result.members().get(id).partitions()); + } + } + + /** + * A, B and C subscribe to T1 and T2, of 4 partitions each, and hold consecutive ranges of both, + * the extra partition going to the first member, as the range assignor assigns them: A holds 0 + * and 1, B holds 2 and C holds 3 of each topic. Each topic has a minimum quota of 1 and 1 extra + * partition, which A keeps for both, so that A has 4 partitions against 2. The balance step + * moves the extra partition of T2 to B: A gives up partition 1 of T2, which B gets. A and B get + * new maps sharing their set of T1, and C keeps its map. + */ + @Test + public void testFromARangeLikeAssignment() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1, T2), Map.of(T1, Set.of(0, 1), T2, Set.of(0, 1))) + .withMember("B", Set.of(T1, T2), Map.of(T1, Set.of(2), T2, Set.of(2))) + .withMember("C", Set.of(T1, T2), Map.of(T1, Set.of(3), T2, Set.of(3))) + .build(); + var describer = new TopicsFixture() + .withTopic(T1, 4) + .withTopic(T2, 4) + .build(); + + var result = assignor.assign(spec, describer); + + assertAssignment(Map.of( + "A", Map.of(T1, Set.of(0, 1), T2, Set.of(0)), + "B", Map.of(T1, Set.of(2), T2, Set.of(1, 2)), + "C", Map.of(T1, Set.of(3), T2, Set.of(3)) + ), result); + } + + /** + * A, B and C subscribe to T1, of 4 partitions, and T2, of 2 partitions, and hold ranges of the + * partitions of both topics taken together, as the uniform assignor can assign them: A holds 0 + * and 1 of T1, B holds 2 and 3 of T1, and C holds both partitions of T2. T1 has a minimum quota + * of 1 and 1 extra partition, which A keeps, as the first of the two owners above the minimum + * quota. T2 has a minimum quota of 0 and 2 extra partitions: C keeps one, and the other goes to + * B, the smallest member without one. B gives up partition 3 of T1 to C, and C gives up + * partition 1 of T2 to B. A keeps its map. + */ + @Test + public void testFromWholeTopics() { Review Comment: It was about C holding all the partitions of T2. I renamed it `testFromATopicHeldByASingleMember`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
