You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by ch...@apache.org on 2026/10/16 17:12:18 UTC
(kafka) branch trunk updated: KAFKA-21214: [6/N] Add the balance step of the uniform2 assignor (#23735)
This is an automated email from the ASF dual-hosted git repository.
chia7712 pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new b35ed0ec911 KAFKA-21214: [6/N] Add the balance step of the uniform2 assignor (#23735)
b35ed0ec911 is described below
commit b35ed0ec911fb6a5a4a31e8481cd81ed5e040bf6
Author: David Jacot <da...@apache.org>
AuthorDate: Fri Oct 16 19:12:10 2026 +0200
KAFKA-21214: [6/N] Add the balance step of the uniform2 assignor (#23735)
uniform2 decides how many partitions of every topic each member gets,
its quotas, in three steps. The first two give every extra partition to
a distinct subscriber of its topic, but they can leave a member with an
assignment much larger than another subscriber of one of its topics.
This patch adds the balance step, `BalanceExtraPartitions`, which the
assignor runs after the other two steps in a later patch. The step
applies a single rule until it no longer applies. The member with the
largest assignment, among those with an extra partition that can move,
gives one to the smallest member that can take it, which is a subscriber
of the topic, with no extra partition of the topic and at least two
partitions smaller. Every move lowers the sum of the squares of the
assignment sizes, so the step ends, and it ends exactly when no extra
partition can move to a subscriber of its topic at least two partitions
smaller.
Changes:
- The givers wait in a heap, the largest first, and a giver without a
receiver is parked until a move may give it one. The step ends once
the largest giver is less than two partitions above the smallest
member.
- The giver gives the extra partition whose move revokes the fewest
partitions. On equal assignment sizes, a giver with a free extra
partition gives first, and a receiver owning more than the minimum
quota of a topic with extra partitions receives first, since only
such a receiver can keep one more of its partitions.
- To find receivers quickly, each giver counts its extra partitions per
cohort, for the topics that the cohort subscribes to. The counts are
kept for every member when they fit, and otherwise for one member at
a time. A constructor forces either way, and a test checks that both
give the same quotas on random groups.
- When all assignment sizes are within one of each other, nothing can
move, and the step returns without building any state.
Reviewers: Chia-Ping Tsai <ch...@gmail.com>
---
.../uniform2/steps/BalanceExtraPartitions.java | 1054 ++++++++++++++++++++
.../uniform2/steps/BalanceExtraPartitionsTest.java | 689 +++++++++++++
2 files changed, 1743 insertions(+)
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/steps/BalanceExtraPartitions.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/steps/BalanceExtraPartitions.java
new file mode 100644
index 00000000000..a548c064db7
--- /dev/null
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/steps/BalanceExtraPartitions.java
@@ -0,0 +1,1054 @@
+/*
+ * 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.steps;
+
+import org.apache.kafka.coordinator.group.assignor.uniform2.CurrentAssignment;
+import org.apache.kafka.coordinator.group.assignor.uniform2.IndexedGroup;
+import org.apache.kafka.coordinator.group.assignor.uniform2.Quotas;
+import org.apache.kafka.coordinator.group.assignor.uniform2.util.IntList;
+import org.apache.kafka.coordinator.group.assignor.uniform2.util.LongHeap;
+
+import java.util.Arrays;
+
+/**
+ * The third step deciding the {@link Quotas}: moves extra partitions until none can move from its
+ * member to another subscriber of its topic whose assignment is at least two partitions smaller.
+ *
+ * <p>An extra partition of a member is owned when the member owns more partitions of its topic
+ * than the minimum quota, so that giving it away revokes one of them. It is free otherwise, and
+ * giving it away revokes no partition.
+ *
+ * <p>The step applies a single rule until it no longer applies. The member with the largest
+ * assignment, among those with an extra partition that can move, gives one to the smallest member
+ * that can take one of its extra partitions. A member can take an extra partition when it
+ * subscribes to the topic, has no extra partition of the topic, and is at least two partitions
+ * smaller. On equal assignment sizes, a giver with a free extra partition gives first, and a
+ * receiver owning more than the minimum quota of a topic with extra partitions receives first,
+ * since only such a receiver can keep one more of its partitions. Member order breaks the
+ * remaining ties. The giver gives the extra partition whose move revokes the fewest partitions. It
+ * tries its free extra partitions first, then its owned ones, and in both lists prefers a topic
+ * that the receiver owns more than the minimum quota of. Ties are broken by taking the last such
+ * topic in the list of the giver.
+ *
+ * <p>Every move goes from a member to one with at least two partitions fewer, so it lowers the sum
+ * of the squares of the assignment sizes by at least two, and the step ends. It ends when no extra
+ * partition can move to a subscriber with at least two partitions fewer, which is the balance
+ * property. With every member subscribed to the same topics, all assignment sizes are then within
+ * one of each other, since a member with two partitions more than another has an extra partition
+ * of a topic the other does not have. Nothing moves when the assignment already has the property.
+ *
+ * <p>Pairing the largest member with the smallest one across cohorts lets a member give an extra
+ * partition to a member of another cohort that needs one. Balancing each cohort on its own could
+ * instead take a move in each of the two cohorts.
+ *
+ * <p>The members with extra partitions wait in a heap, the largest first. A member without a
+ * receiver is parked until a move may give it one, see {@link #wake}. Receiving an extra partition
+ * also wakes a member, since it is then larger and has a new topic to give. The step ends when the
+ * largest giver is less than two partitions above the smallest member.
+ *
+ * <p>Receivers are searched cohort by cohort, starting with the cohort of the smallest member, and
+ * within a cohort from its smallest member. A cohort is skipped when each of its members already
+ * has an extra partition of every topic of the giver that the cohort subscribes to. A member is
+ * known to be able to take an extra partition of the giver when the member has fewer extra
+ * partitions in all than the giver has of topics of the member's cohort. Otherwise, their topics
+ * are compared.
+ */
+public final class BalanceExtraPartitions {
+ /**
+ * The offset of the free extra partitions in {@link #resumePositions}.
+ */
+ private static final int FREE = 0;
+
+ /**
+ * The offset of the owned extra partitions in {@link #resumePositions}.
+ */
+ private static final int OWNED = 2;
+
+ /**
+ * The members and topics of the group.
+ */
+ private final IndexedGroup group;
+
+ /**
+ * The current assignment, which tells which partitions a move revokes.
+ */
+ private final CurrentAssignment currentAssignment;
+
+ /**
+ * The quotas, which the step updates.
+ */
+ private final Quotas quotas;
+
+ /**
+ * Whether the counts per cohort of every member are kept, in the fields ending in {@code Of},
+ * see {@link #countsFit}. Otherwise, only those of {@link #countedMember} are, in the fields
+ * starting with {@code counted}. The counts per cohort of a member are, per cohort, the number
+ * of its extra partitions of topics that the cohort subscribes to, and the number of those
+ * topics that are full in the cohort, see {@link #membersWithExtraPartitionPerCohort}.
+ */
+ private final boolean keepCounts;
+
+ /**
+ * Per member, the topics of its free extra partitions.
+ */
+ private IntList[] freeExtraPartitions;
+
+ /**
+ * Per member, the topics of its owned extra partitions.
+ */
+ private IntList[] ownedExtraPartitions;
+
+ /**
+ * Per member, the topics with extra partitions of which it owns more than the minimum quota.
+ */
+ private IntList[] ownedAboveMinQuota;
+
+ /**
+ * Per topic, {@link #ownedMark} when {@link #ownedMarkedMember} owns more than the minimum
+ * quota of the topic.
+ */
+ private int[] ownedMarks;
+
+ /**
+ * The current mark of {@link #ownedMarks}.
+ */
+ private int ownedMark;
+
+ /**
+ * The member whose entries in {@link #ownedMarks} are current, or -1.
+ */
+ private int ownedMarkedMember;
+
+ /**
+ * Per topic with extra partitions, per cohort subscribing to it in the order of
+ * {@link IndexedGroup#cohortsOf}, the number of members of the cohort with an extra partition
+ * of the topic. The entry of a topic without extra partitions is null. The topic is full in the
+ * cohort when every member of the cohort has an extra partition of it, and then no member of
+ * the cohort can take another.
+ */
+ private int[][] membersWithExtraPartitionPerCohort;
+
+ /**
+ * Per cohort, its number of members.
+ */
+ private int[] cohortSizes;
+
+ /**
+ * The members that may give, the largest first, see {@link #giverKey}. The heap holds each
+ * member with the key it had when it was added. A member whose key changes is added again, and
+ * the stale entries are skipped.
+ */
+ private LongHeap givers;
+
+ /**
+ * Per cohort, its members, the smallest first, see {@link #receiverKey}, with the same kind of
+ * entries as {@link #givers}.
+ */
+ private LongHeap[] receivers;
+
+ /**
+ * The number of members of each assignment size, at the index of the assignment size minus
+ * {@link #initialSmallestSize}. A move from a member of size {@code L} to one of size
+ * {@code R <= L - 2} leaves them at {@code L - 1} and {@code R + 1}, so the assignment sizes
+ * stay within their range at the start.
+ */
+ private int[] sizeCounts;
+
+ /**
+ * The smallest assignment size at the start.
+ */
+ private int initialSmallestSize;
+
+ /**
+ * The smallest assignment size of any member. It never decreases.
+ */
+ private int smallestAssignmentSize;
+
+ /**
+ * Per member, whether it is parked. A member is parked when no member can take any of its
+ * extra partitions, see {@link #wake}.
+ */
+ private boolean[] isParked;
+
+ /**
+ * The parked members by assignment size, at the index of the assignment size minus
+ * {@link #initialSmallestSize}, created when the first member is parked. A member is listed
+ * under the assignment size it had when it was parked, and the entry only holds while the
+ * member stays parked at that size.
+ */
+ private IntList[] parkedBySize;
+
+ /**
+ * When the counts per cohort are kept, per member, per cohort, the number of its extra
+ * partitions of topics that the cohort subscribes to. The entry of a member is null until a
+ * receiver is first searched for it.
+ */
+ private int[][] extraPartitionsPerCohortOf;
+
+ /**
+ * When the counts per cohort are kept, per member, per cohort, the number of the topics of its
+ * extra partitions that are full in the cohort. The entry of a member is null until a receiver
+ * is first searched for it. When it equals the entry of {@link #extraPartitionsPerCohortOf}, no
+ * member of the cohort can take an extra partition of the member.
+ */
+ private int[][] fullTopicsPerCohortOf;
+
+ /**
+ * When the counts per cohort are kept, per member, the cohorts subscribing to a topic of its
+ * extra partitions. The entry of a member is null until a receiver is first searched for it.
+ * It may also hold cohorts whose count dropped to zero since the last search, and may hold a
+ * cohort more than once.
+ */
+ private IntList[] subscribingCohortsOf;
+
+ /**
+ * When the counts per cohort are not kept, the number of extra partitions per cohort of
+ * {@link #countedMember}, see {@link #extraPartitionsPerCohortOf}.
+ */
+ private int[] countedExtraPartitionsPerCohort;
+
+ /**
+ * When the counts per cohort are not kept, the number of full topics per cohort of
+ * {@link #countedMember}, see {@link #fullTopicsPerCohortOf}.
+ */
+ private int[] countedFullTopicsPerCohort;
+
+ /**
+ * When the counts per cohort are not kept, the cohorts subscribing to a topic of the extra
+ * partitions of {@link #countedMember}. It may also hold cohorts whose count dropped to zero
+ * since the last search.
+ */
+ private IntList countedSubscribingCohorts;
+
+ /**
+ * When the counts per cohort are not kept, the member whose counts are held, or -1.
+ */
+ private int countedMember;
+
+ /**
+ * Per topic, {@link #mark} when {@link #markedMember} has an extra partition of the topic.
+ */
+ private int[] marks;
+
+ /**
+ * The current mark. A topic is marked when its entry in {@link #marks} holds it, so that
+ * marking another member needs no clearing.
+ */
+ private int mark;
+
+ /**
+ * The member whose entries in {@link #marks} are current, or -1.
+ */
+ private int markedMember;
+
+ /**
+ * Per member as a giver, the receiver of its last move, or -1. For that receiver, the topics
+ * after the positions of {@link #resumePositions} in the lists of the giver were already looked
+ * at. They stay so while the receiver gives nothing and the giver receives nothing, since
+ * removing a topic from a list moves the last topic, already looked at, into its slot.
+ */
+ private int[] lastReceiver;
+
+ /**
+ * Per member as a giver, the number of extra partitions that {@link #lastReceiver} had given
+ * when the member last gave to it.
+ */
+ private int[] lastReceiverGiven;
+
+ /**
+ * Per member, the number of extra partitions it gave.
+ */
+ private int[] given;
+
+ /**
+ * Per member as a giver, the positions in its free ({@link #FREE}) and owned ({@link #OWNED})
+ * extra partitions from which the next move to {@link #lastReceiver} searches for a topic. The
+ * position at {@code 4 * member + list} is for a topic the receiver can take, and the one at
+ * {@code 4 * member + list + 1} for a topic the receiver can take and owns more than the
+ * minimum quota of. The topics after a position were already looked at, and the moves do not
+ * change why they were passed over.
+ */
+ private int[] resumePositions;
+
+ /**
+ * The members taken out of a heap of receivers while searching it for a receiver. They are put
+ * back after the search.
+ */
+ private IntList skipped;
+
+ /**
+ * The receiver keys of the smallest member of each cohort searched for a receiver, the
+ * smallest first, see {@link #receiverKey}.
+ */
+ private LongHeap smallestMembersOfCohorts;
+
+ /**
+ * Creates the step, which keeps the counts per cohort of every member when they fit, see
+ * {@link #countsFit}.
+ *
+ * @param group The members and topics of the group.
+ * @param currentAssignment The current assignment.
+ * @param quotas The quotas to balance, as the previous steps left them.
+ */
+ public BalanceExtraPartitions(IndexedGroup group, CurrentAssignment currentAssignment, Quotas quotas) {
+ this(group, currentAssignment, quotas, countsFit(group));
+ }
+
+ /**
+ * Creates the step, keeping the counts per cohort of every member or of one member at a time.
+ * Tests use it to run the step both ways on the same group.
+ *
+ * @param group The members and topics of the group.
+ * @param currentAssignment The current assignment.
+ * @param quotas The quotas to balance, as the previous steps left them.
+ * @param keepCounts Whether the counts per cohort of every member are kept.
+ */
+ public BalanceExtraPartitions(
+ IndexedGroup group,
+ CurrentAssignment currentAssignment,
+ Quotas quotas,
+ boolean keepCounts
+ ) {
+ this.group = group;
+ this.currentAssignment = currentAssignment;
+ this.quotas = quotas;
+ this.keepCounts = keepCounts;
+ }
+
+ /**
+ * Tells whether the counts per cohort of every member fit, which is when they take no more
+ * entries than the subscriptions of the members, one per member and subscribed topic. The
+ * counts take one entry per member and cohort.
+ *
+ * @return {@code true} if the number of members times the number of cohorts is at most the sum
+ * of the sizes of the subscriptions of the members.
+ */
+ static boolean countsFit(IndexedGroup group) {
+ long subscribedTopics = 0;
+ for (int cohort = 0; cohort < group.cohortCount(); cohort++) {
+ subscribedTopics += (long) group.membersOf(cohort).length * group.topicsOf(cohort).length;
+ }
+ return (long) group.memberCount() * group.cohortCount() <= subscribedTopics;
+ }
+
+ /**
+ * Balances the quotas, see the class documentation. When all assignment sizes are within one
+ * of each other, nothing can move, and it returns at once without building any state.
+ */
+ public void run() {
+ if (!mayMove(group, quotas)) {
+ return;
+ }
+ createState();
+ balance();
+ }
+
+ /**
+ * Tells whether an extra partition may move, which needs a member at least two partitions
+ * larger than another. With every member subscribed to the same topics, an extra partition can
+ * then always move. Otherwise it may not, for example when the two members subscribe to
+ * different topics. The step then builds its state and moves nothing.
+ *
+ * @return {@code false} when all assignment sizes are within one of each other, in which case
+ * nothing can move.
+ */
+ static boolean mayMove(IndexedGroup group, Quotas quotas) {
+ int smallest = Integer.MAX_VALUE;
+ // Assignment sizes are not negative, so an empty group returns false.
+ int largest = 0;
+ for (int member = 0; member < group.memberCount(); member++) {
+ smallest = Math.min(smallest, quotas.assignmentSize(member));
+ largest = Math.max(largest, quotas.assignmentSize(member));
+ }
+ return largest - smallest >= 2;
+ }
+
+ /**
+ * Creates the state of the step from the quotas as the previous steps left them.
+ */
+ private void createState() {
+ int memberCount = group.memberCount();
+ cohortSizes = new int[group.cohortCount()];
+ for (int cohort = 0; cohort < group.cohortCount(); cohort++) {
+ cohortSizes[cohort] = group.membersOf(cohort).length;
+ }
+ freeExtraPartitions = new IntList[memberCount];
+ ownedExtraPartitions = new IntList[memberCount];
+ ownedAboveMinQuota = new IntList[memberCount];
+ for (int member = 0; member < memberCount; member++) {
+ freeExtraPartitions[member] = new IntList(quotas.extraPartitionCount(member));
+ ownedExtraPartitions[member] = new IntList(0);
+ ownedAboveMinQuota[member] = new IntList(0);
+ }
+ membersWithExtraPartitionPerCohort = new int[group.topicCount()][];
+ recordExtraPartitions();
+ createHeaps();
+ computeSizeCounts();
+ isParked = new boolean[memberCount];
+ if (keepCounts) {
+ extraPartitionsPerCohortOf = new int[memberCount][];
+ fullTopicsPerCohortOf = new int[memberCount][];
+ subscribingCohortsOf = new IntList[memberCount];
+ } else {
+ countedExtraPartitionsPerCohort = new int[group.cohortCount()];
+ countedFullTopicsPerCohort = new int[group.cohortCount()];
+ countedSubscribingCohorts = new IntList(16);
+ countedMember = -1;
+ }
+ marks = new int[group.topicCount()];
+ markedMember = -1;
+ ownedMarks = new int[group.topicCount()];
+ ownedMarkedMember = -1;
+ lastReceiver = new int[memberCount];
+ Arrays.fill(lastReceiver, -1);
+ lastReceiverGiven = new int[memberCount];
+ given = new int[memberCount];
+ resumePositions = new int[4 * memberCount];
+ skipped = new IntList(16);
+ smallestMembersOfCohorts = new LongHeap(group.cohortCount());
+ }
+
+ /**
+ * Records the free and owned extra partitions of every member,
+ * {@link #membersWithExtraPartitionPerCohort}, and the topics with extra partitions that each
+ * member owns more than the minimum quota of.
+ */
+ private void recordExtraPartitions() {
+ // Per member, the number of partitions of the topic that it owns.
+ var ownedPartitionCounts = new int[group.memberCount()];
+ var owners = new IntList(16);
+ for (int topic = 0; topic < group.topicCount(); topic++) {
+ if (quotas.extraPartitions(topic) == 0) {
+ continue;
+ }
+ currentAssignment.countOwned(topic, ownedPartitionCounts, owners);
+ membersWithExtraPartitionPerCohort[topic] = new int[group.cohortsOf(topic).length];
+ var members = quotas.membersWithExtraPartition(topic);
+ for (int i = 0; i < members.size(); i++) {
+ int member = members.get(i);
+ boolean owned = ownedPartitionCounts[member] > quotas.minQuota(topic);
+ (owned ? ownedExtraPartitions : freeExtraPartitions)[member].add(topic);
+ membersWithExtraPartitionPerCohort[topic][cohortIndex(topic, group.cohortOf(member))]++;
+ }
+ for (int i = 0; i < owners.size(); i++) {
+ int owner = owners.get(i);
+ if (ownedPartitionCounts[owner] > quotas.minQuota(topic)) {
+ ownedAboveMinQuota[owner].add(topic);
+ }
+ // Reset for the next topic, as countOwned requires.
+ ownedPartitionCounts[owner] = 0;
+ }
+ }
+ }
+
+ /**
+ * Creates the heap of the givers, the members with extra partitions, and the heaps of the
+ * receivers, every cohort with all its members.
+ */
+ private void createHeaps() {
+ var giverKeys = new long[group.memberCount()];
+ int giverCount = 0;
+ receivers = new LongHeap[group.cohortCount()];
+ for (int cohort = 0; cohort < group.cohortCount(); cohort++) {
+ var members = group.membersOf(cohort);
+ // The heap keeps this array as its storage. Every move adds its giver and its receiver
+ // again to the heaps of their cohorts, so the array has room for as many entries again.
+ var receiverKeys = new long[2 * members.length];
+ for (int i = 0; i < members.length; i++) {
+ receiverKeys[i] = receiverKey(members[i]);
+ if (quotas.extraPartitionCount(members[i]) > 0) {
+ giverKeys[giverCount++] = giverKey(members[i]);
+ }
+ }
+ receivers[cohort] = new LongHeap(receiverKeys, members.length);
+ }
+ givers = new LongHeap(giverKeys, giverCount);
+ }
+
+ /**
+ * Computes {@link #sizeCounts} and the smallest assignment size at the start.
+ */
+ private void computeSizeCounts() {
+ int min = Integer.MAX_VALUE;
+ int max = Integer.MIN_VALUE;
+ for (int member = 0; member < group.memberCount(); member++) {
+ min = Math.min(min, quotas.assignmentSize(member));
+ max = Math.max(max, quotas.assignmentSize(member));
+ }
+ initialSmallestSize = min;
+ smallestAssignmentSize = min;
+ sizeCounts = new int[max - min + 1];
+ for (int member = 0; member < group.memberCount(); member++) {
+ sizeCounts[quotas.assignmentSize(member) - min]++;
+ }
+ }
+
+ /**
+ * Moves extra partitions from the largest giver to its receiver until no member can give.
+ */
+ private void balance() {
+ while (!givers.isEmpty()) {
+ long key = givers.poll();
+ int giver = member(key);
+ if (key != giverKey(giver) || isParked[giver] || quotas.extraPartitionCount(giver) == 0) {
+ continue;
+ }
+ // The other givers are no larger, and the parked members have no receiver either. Once
+ // the largest giver is less than two partitions above the smallest member, no member
+ // can give any more.
+ if (quotas.assignmentSize(giver) - 2 < smallestAssignmentSize) {
+ return;
+ }
+ int receiver = findReceiver(giver);
+ if (receiver < 0) {
+ park(giver);
+ continue;
+ }
+ move(giver, receiver);
+ // The receiver has a new extra partition to give, and both members changed size.
+ isParked[receiver] = false;
+ givers.add(giverKey(receiver));
+ receivers[group.cohortOf(receiver)].add(receiverKey(receiver));
+ if (quotas.extraPartitionCount(giver) > 0) {
+ givers.add(giverKey(giver));
+ }
+ receivers[group.cohortOf(giver)].add(receiverKey(giver));
+ wake(giver);
+ }
+ }
+
+ /**
+ * Finds the receiver of an extra partition of the giver. It is the first member in the order of
+ * {@link #receiverKey} among the members at least two partitions smaller that can take an
+ * extra partition of the giver. The cohorts are searched from the one with the smallest member,
+ * so that the members of the other cohorts are mostly compared by assignment size only.
+ *
+ * @return The receiver, or -1 if there is none.
+ */
+ private int findReceiver(int giver) {
+ // The largest assignment size that a receiver may have.
+ int maxSize = quotas.assignmentSize(giver) - 2;
+ // Per cohort, the number of extra partitions of the giver of topics that the cohort
+ // subscribes to, and the number of those topics that are full in the cohort.
+ int[] extraPartitions = extraPartitionsPerCohort(giver);
+ int[] fullTopics = fullTopicsPerCohort(giver);
+ var cohorts = subscribingCohorts(giver);
+ smallestMembersOfCohorts.clear();
+ int kept = 0;
+ for (int i = 0; i < cohorts.size(); i++) {
+ int cohort = cohorts.get(i);
+ // The cohorts no longer subscribing to a topic of the giver are dropped on the way.
+ if (extraPartitions[cohort] == 0) {
+ continue;
+ }
+ cohorts.set(kept++, cohort);
+ // In a cohort where every topic of the giver is full, no member can take one.
+ if (fullTopics[cohort] == extraPartitions[cohort]) {
+ continue;
+ }
+ long key = peekSmallestKey(cohort);
+ if (key >= 0 && quotas.assignmentSize(member(key)) <= maxSize) {
+ smallestMembersOfCohorts.add(key);
+ }
+ }
+ cohorts.truncate(kept);
+ // The first receiver found so far in the order of receiverKey, or -1.
+ int receiver = -1;
+ while (!smallestMembersOfCohorts.isEmpty()) {
+ long key = smallestMembersOfCohorts.poll();
+ if (receiver >= 0 && key > receiverKey(receiver)) {
+ break;
+ }
+ int cohortReceiver = findReceiverInCohort(giver, group.cohortOf(member(key)), maxSize, receiver);
+ if (cohortReceiver >= 0) {
+ receiver = cohortReceiver;
+ }
+ }
+ return receiver;
+ }
+
+ /**
+ * Drops the stale entries at the top of the heap of receivers of the cohort.
+ *
+ * @return The key of the first member of the cohort in the order of {@link #receiverKey}, or -1
+ * if there is none.
+ */
+ private long peekSmallestKey(int cohort) {
+ var heap = receivers[cohort];
+ while (!heap.isEmpty()) {
+ long key = heap.peek();
+ if (key == receiverKey(member(key))) {
+ return key;
+ }
+ heap.poll();
+ }
+ return -1;
+ }
+
+ /**
+ * Finds the first member of the cohort in the order of {@link #receiverKey} that can take an
+ * extra partition of the giver. The member must have at most {@code maxSize} partitions and,
+ * if there is a best receiver so far, come before it.
+ *
+ * @return The member, or -1 if there is none.
+ */
+ private int findReceiverInCohort(int giver, int cohort, int maxSize, int bestReceiver) {
+ var heap = receivers[cohort];
+ int receiver = -1;
+ long key;
+ while ((key = peekSmallestKey(cohort)) >= 0) {
+ int member = member(key);
+ if (quotas.assignmentSize(member) > maxSize || bestReceiver >= 0 && key > receiverKey(bestReceiver)) {
+ break;
+ }
+ if (canTake(giver, member)) {
+ receiver = member;
+ break;
+ }
+ skipped.add(member);
+ heap.poll();
+ }
+ for (int i = 0; i < skipped.size(); i++) {
+ heap.add(receiverKey(skipped.get(i)));
+ }
+ skipped.clear();
+ return receiver;
+ }
+
+ /**
+ * @return {@code true} if the receiver subscribes to a topic of an extra partition of the giver
+ * and has no extra partition of it.
+ */
+ private boolean canTake(int giver, int receiver) {
+ // The receiver only has extra partitions of topics of its cohort. When the giver has more
+ // of these topics than the receiver has extra partitions, the receiver lacks one of them.
+ if ((keepCounts || giver == countedMember)
+ && extraPartitionsPerCohort(giver)[group.cohortOf(receiver)] > quotas.extraPartitionCount(receiver)) {
+ return true;
+ }
+ markExtraPartitions(receiver);
+ return canTakeOne(freeExtraPartitions[giver], receiver) || canTakeOne(ownedExtraPartitions[giver], receiver);
+ }
+
+ /**
+ * Tells whether the receiver can take an extra partition of one of the topics. The extra
+ * partitions of the receiver must be marked.
+ *
+ * @return {@code true} if the receiver subscribes to one of the topics and has no extra
+ * partition of it.
+ */
+ private boolean canTakeOne(IntList topics, int receiver) {
+ for (int i = 0; i < topics.size(); i++) {
+ int topic = topics.get(i);
+ if (marks[topic] != mark && subscribes(receiver, topic)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /**
+ * Moves an extra partition from the giver to the receiver. The receiver must be able to take
+ * an extra partition of the giver. The giver gives a free extra partition if the receiver can
+ * take one, and an owned one otherwise. In both lists, it prefers a topic that the receiver
+ * owns more than the minimum quota of, which revokes one partition fewer from the receiver.
+ */
+ private void move(int giver, int receiver) {
+ markExtraPartitions(receiver);
+ markOwnedAboveMinQuota(receiver);
+ // The topics already looked at for the same receiver need not be looked at again.
+ if (lastReceiver[giver] != receiver || lastReceiverGiven[giver] != given[receiver]) {
+ lastReceiver[giver] = receiver;
+ lastReceiverGiven[giver] = given[receiver];
+ Arrays.fill(resumePositions, 4 * giver, 4 * giver + 4, Integer.MAX_VALUE);
+ }
+ var topics = freeExtraPartitions[giver];
+ int chosen = chooseTopicIndex(topics, receiver, 4 * giver + FREE);
+ if (chosen < 0) {
+ topics = ownedExtraPartitions[giver];
+ chosen = chooseTopicIndex(topics, receiver, 4 * giver + OWNED);
+ }
+ int topic = topics.get(chosen);
+ updateCountsForMove(giver, topic, receiver);
+ quotas.moveExtraPartition(giver, topic, receiver);
+ topics.remove(chosen);
+ given[giver]++;
+ // The extra partitions of the receiver grow, so what it looked at as a giver no longer
+ // holds.
+ lastReceiver[receiver] = -1;
+ (ownsAboveMinQuota(topic) ? ownedExtraPartitions : freeExtraPartitions)[receiver].add(topic);
+ // The receiver is the marked member, so its marks stay exact.
+ marks[topic] = mark;
+ updateSizeCounts(giver, receiver);
+ }
+
+ /**
+ * Chooses the topic to move among the topics, the free or the owned extra partitions of the
+ * giver. It prefers a topic that the receiver owns more than the minimum quota of, and breaks
+ * ties by taking the last one in the list. The extra partitions of the receiver and its topics
+ * owned more than the minimum quota of must be marked. The search starts from the resume
+ * positions at the position, which it updates.
+ *
+ * @return The index of the chosen topic in the topics, or -1 if the receiver can take none.
+ */
+ private int chooseTopicIndex(IntList topics, int receiver, int position) {
+ if (!ownedAboveMinQuota[receiver].isEmpty()) {
+ int chosen = findLastTopicIndex(topics, receiver, resumePositions[position + 1], true);
+ // The topic moved into the slot of the chosen one, if any, was already looked at.
+ resumePositions[position + 1] = chosen < 0 ? -1 : chosen - 1;
+ if (chosen >= 0) {
+ return chosen;
+ }
+ }
+ int chosen = findLastTopicIndex(topics, receiver, resumePositions[position], false);
+ resumePositions[position] = chosen < 0 ? -1 : chosen - 1;
+ return chosen;
+ }
+
+ /**
+ * Finds the last topic, at most at the index {@code from}, that the receiver can take and, if
+ * {@code ownedByReceiver}, owns more than the minimum quota of. The extra partitions of the
+ * receiver must be marked, and so must its topics owned more than the minimum quota of if
+ * {@code ownedByReceiver}.
+ *
+ * @return The index of the topic in the topics, or -1 if there is none.
+ */
+ private int findLastTopicIndex(IntList topics, int receiver, int from, boolean ownedByReceiver) {
+ for (int i = Math.min(topics.size() - 1, from); i >= 0; i--) {
+ int topic = topics.get(i);
+ if (marks[topic] != mark && subscribes(receiver, topic) && (!ownedByReceiver || ownsAboveMinQuota(topic))) {
+ return i;
+ }
+ }
+ return -1;
+ }
+
+ /**
+ * Updates {@link #membersWithExtraPartitionPerCohort} and the counts per cohort for the move of
+ * an extra partition of the topic from the giver to the receiver. It must be called before the
+ * move.
+ */
+ private void updateCountsForMove(int giver, int topic, int receiver) {
+ int[] cohorts = group.cohortsOf(topic);
+ // The search for a receiver counted the giver just before the move. The giver loses the
+ // topic, as it was before the move.
+ int[] extraPartitions = extraPartitionsPerCohort(giver);
+ int[] fullTopics = fullTopicsPerCohort(giver);
+ for (int index = 0; index < cohorts.length; index++) {
+ extraPartitions[cohorts[index]]--;
+ fullTopics[cohorts[index]] -= isFull(topic, index) ? 1 : 0;
+ }
+ updateMembersWithExtraPartitionPerCohort(giver, topic, receiver);
+ // The receiver gets the topic, as it is after the move.
+ if (keepCounts && extraPartitionsPerCohortOf[receiver] != null) {
+ for (int index = 0; index < cohorts.length; index++) {
+ if (extraPartitionsPerCohortOf[receiver][cohorts[index]]++ == 0) {
+ subscribingCohortsOf[receiver].add(cohorts[index]);
+ }
+ fullTopicsPerCohortOf[receiver][cohorts[index]] += isFull(topic, index) ? 1 : 0;
+ }
+ }
+ }
+
+ /**
+ * Moves a member with an extra partition of the topic from the cohort of the giver to the
+ * cohort of the receiver, in {@link #membersWithExtraPartitionPerCohort}. When the topic stops
+ * being full in the first cohort or becomes full in the second, the counts of full topics kept
+ * for the other members with an extra partition of the topic are updated.
+ */
+ private void updateMembersWithExtraPartitionPerCohort(int giver, int topic, int receiver) {
+ int giverIndex = cohortIndex(topic, group.cohortOf(giver));
+ int receiverIndex = cohortIndex(topic, group.cohortOf(receiver));
+ if (giverIndex == receiverIndex) {
+ return;
+ }
+ boolean wasFull = isFull(topic, giverIndex);
+ membersWithExtraPartitionPerCohort[topic][giverIndex]--;
+ membersWithExtraPartitionPerCohort[topic][receiverIndex]++;
+ boolean becomesFull = isFull(topic, receiverIndex);
+ if (!keepCounts || !wasFull && !becomesFull) {
+ return;
+ }
+ int giverCohort = group.cohortOf(giver);
+ int receiverCohort = group.cohortOf(receiver);
+ var members = quotas.membersWithExtraPartition(topic);
+ for (int i = 0; i < members.size(); i++) {
+ int member = members.get(i);
+ if (member != giver && fullTopicsPerCohortOf[member] != null) {
+ fullTopicsPerCohortOf[member][giverCohort] -= wasFull ? 1 : 0;
+ fullTopicsPerCohortOf[member][receiverCohort] += becomesFull ? 1 : 0;
+ }
+ }
+ }
+
+ /**
+ * @return {@code true} if every member of the cohort at the index of the cohorts of the topic
+ * has an extra partition of the topic.
+ */
+ private boolean isFull(int topic, int index) {
+ return membersWithExtraPartitionPerCohort[topic][index] == cohortSizes[group.cohortsOf(topic)[index]];
+ }
+
+ /**
+ * Returns the number of full topics per cohort of the member, see
+ * {@link #fullTopicsPerCohortOf}. It must be called after {@link #extraPartitionsPerCohort} for
+ * the member.
+ */
+ private int[] fullTopicsPerCohort(int member) {
+ return keepCounts ? fullTopicsPerCohortOf[member] : countedFullTopicsPerCohort;
+ }
+
+ /**
+ * Returns the number of extra partitions per cohort of the member, see
+ * {@link #extraPartitionsPerCohortOf}. The counts per cohort of the member are computed when
+ * asked for the first time or, when they are not kept, when asked for another member than the
+ * last one.
+ */
+ private int[] extraPartitionsPerCohort(int member) {
+ if (keepCounts) {
+ if (extraPartitionsPerCohortOf[member] == null) {
+ extraPartitionsPerCohortOf[member] = new int[group.cohortCount()];
+ fullTopicsPerCohortOf[member] = new int[group.cohortCount()];
+ subscribingCohortsOf[member] = new IntList(4);
+ computeCounts(
+ member,
+ extraPartitionsPerCohortOf[member],
+ fullTopicsPerCohortOf[member],
+ subscribingCohortsOf[member]
+ );
+ }
+ return extraPartitionsPerCohortOf[member];
+ }
+ if (member != countedMember) {
+ for (int i = 0; i < countedSubscribingCohorts.size(); i++) {
+ countedExtraPartitionsPerCohort[countedSubscribingCohorts.get(i)] = 0;
+ countedFullTopicsPerCohort[countedSubscribingCohorts.get(i)] = 0;
+ }
+ countedSubscribingCohorts.clear();
+ computeCounts(member, countedExtraPartitionsPerCohort, countedFullTopicsPerCohort, countedSubscribingCohorts);
+ countedMember = member;
+ }
+ return countedExtraPartitionsPerCohort;
+ }
+
+ /**
+ * Returns the cohorts subscribing to a topic of the extra partitions of the member, see
+ * {@link #subscribingCohortsOf}. It must be called after {@link #extraPartitionsPerCohort} for
+ * the member.
+ */
+ private IntList subscribingCohorts(int member) {
+ return keepCounts ? subscribingCohortsOf[member] : countedSubscribingCohorts;
+ }
+
+ /**
+ * Computes the counts per cohort of the member, see {@link #keepCounts}. The counts must be
+ * zeroed and the list of cohorts empty.
+ */
+ private void computeCounts(int member, int[] extraPartitions, int[] fullTopics, IntList cohorts) {
+ addToCounts(freeExtraPartitions[member], extraPartitions, fullTopics, cohorts);
+ addToCounts(ownedExtraPartitions[member], extraPartitions, fullTopics, cohorts);
+ }
+
+ /**
+ * Adds the topics to the counts per cohort, and the cohorts they reach for the first time to
+ * the list of cohorts.
+ */
+ private void addToCounts(IntList topics, int[] extraPartitions, int[] fullTopics, IntList cohorts) {
+ for (int i = 0; i < topics.size(); i++) {
+ int topic = topics.get(i);
+ int[] topicCohorts = group.cohortsOf(topic);
+ int[] membersPerCohort = membersWithExtraPartitionPerCohort[topic];
+ for (int index = 0; index < topicCohorts.length; index++) {
+ int cohort = topicCohorts[index];
+ if (extraPartitions[cohort]++ == 0) {
+ cohorts.add(cohort);
+ }
+ fullTopics[cohort] += membersPerCohort[index] == cohortSizes[cohort] ? 1 : 0;
+ }
+ }
+ }
+
+ /**
+ * Updates {@link #sizeCounts} and {@link #smallestAssignmentSize} after a move from the giver,
+ * now one partition smaller, to the receiver, now one partition larger. It must be called after
+ * the move.
+ */
+ private void updateSizeCounts(int giver, int receiver) {
+ int giverIndex = quotas.assignmentSize(giver) - initialSmallestSize;
+ int receiverIndex = quotas.assignmentSize(receiver) - initialSmallestSize;
+ sizeCounts[giverIndex + 1]--;
+ sizeCounts[giverIndex]++;
+ sizeCounts[receiverIndex - 1]--;
+ sizeCounts[receiverIndex]++;
+ while (sizeCounts[smallestAssignmentSize - initialSmallestSize] == 0) {
+ smallestAssignmentSize++;
+ }
+ }
+
+ /**
+ * Parks the giver when no receiver is found for it. A giver at most two partitions above the
+ * smallest assignment size is not listed, since no move can wake it. A move leaves its giver
+ * above the smallest assignment size, which never decreases, and only wakes members two
+ * partitions above its giver, see {@link #wake}.
+ */
+ private void park(int giver) {
+ isParked[giver] = true;
+ if (quotas.assignmentSize(giver) - 2 <= smallestAssignmentSize) {
+ return;
+ }
+ if (parkedBySize == null) {
+ parkedBySize = new IntList[sizeCounts.length];
+ }
+ int index = quotas.assignmentSize(giver) - initialSmallestSize;
+ if (parkedBySize[index] == null) {
+ parkedBySize[index] = new IntList(4);
+ }
+ parkedBySize[index].add(giver);
+ }
+
+ /**
+ * Wakes the parked members that can now give to the giver of the last move. These are the
+ * members exactly two partitions larger than the giver from which the giver can take an extra
+ * partition.
+ *
+ * <p>A parked member has no receiver, and only a move can give it one. The receiver of the move
+ * gets larger and gets an extra partition, which makes it no better a receiver. The giver gets
+ * one partition smaller, at an assignment size of {@code s}, and lacks the topic it gave. A
+ * parked member can then give to it only through:
+ * <ul>
+ * <li>the topic it gave, if the parked member has an extra partition of it and has at least
+ * {@code s + 2} partitions. This never happens, since the receiver of the move lacked the
+ * topic and had at most {@code s - 1} partitions, so it was already a receiver of the
+ * parked member;</li>
+ * <li>another topic that the giver lacks. The giver lacked it before the move too, when it
+ * had {@code s + 1} partitions, and was then no receiver of the parked member. So the
+ * parked member has at most {@code s + 2} partitions, and exactly {@code s + 2} if the
+ * giver is now its receiver.</li>
+ * </ul>
+ */
+ private void wake(int lastGiver) {
+ int size = quotas.assignmentSize(lastGiver);
+ int index = size + 2 - initialSmallestSize;
+ if (parkedBySize == null || index >= parkedBySize.length || parkedBySize[index] == null) {
+ return;
+ }
+ var parked = parkedBySize[index];
+ int kept = 0;
+ for (int i = 0; i < parked.size(); i++) {
+ int parkedMember = parked.get(i);
+ if (!isParked[parkedMember] || quotas.assignmentSize(parkedMember) != size + 2) {
+ continue;
+ }
+ if (canTake(parkedMember, lastGiver)) {
+ isParked[parkedMember] = false;
+ givers.add(giverKey(parkedMember));
+ } else {
+ parked.set(kept++, parkedMember);
+ }
+ }
+ parked.truncate(kept);
+ }
+
+ /**
+ * Marks the topics of the extra partitions of the member, unless they are marked already.
+ */
+ private void markExtraPartitions(int member) {
+ if (member == markedMember) {
+ return;
+ }
+ mark++;
+ markTopics(freeExtraPartitions[member]);
+ markTopics(ownedExtraPartitions[member]);
+ markedMember = member;
+ }
+
+ /**
+ * Marks the topics with the current mark.
+ */
+ private void markTopics(IntList topics) {
+ for (int i = 0; i < topics.size(); i++) {
+ marks[topics.get(i)] = mark;
+ }
+ }
+
+ /**
+ * Marks the topics with extra partitions of which the member owns more than the minimum quota,
+ * unless they are marked already.
+ */
+ private void markOwnedAboveMinQuota(int member) {
+ if (member == ownedMarkedMember) {
+ return;
+ }
+ ownedMark++;
+ var topics = ownedAboveMinQuota[member];
+ for (int i = 0; i < topics.size(); i++) {
+ ownedMarks[topics.get(i)] = ownedMark;
+ }
+ ownedMarkedMember = member;
+ }
+
+ /**
+ * @return {@code true} if the member marked by {@link #markOwnedAboveMinQuota} owns more than
+ * the minimum quota of the topic.
+ */
+ private boolean ownsAboveMinQuota(int topic) {
+ return ownedMarks[topic] == ownedMark;
+ }
+
+ /**
+ * @return {@code true} if the member subscribes to the topic.
+ */
+ private boolean subscribes(int member, int topic) {
+ // A topic usually has far fewer subscribing cohorts than a cohort has topics, so the
+ // cohorts of the topic are searched.
+ return Arrays.binarySearch(group.cohortsOf(topic), group.cohortOf(member)) >= 0;
+ }
+
+ /**
+ * @return The index of the cohort among the cohorts subscribing to the topic.
+ */
+ private int cohortIndex(int topic, int cohort) {
+ return Arrays.binarySearch(group.cohortsOf(topic), cohort);
+ }
+
+ /**
+ * Returns the key of the member in the min heap of givers. The high 32 bits hold its
+ * assignment size, complemented so that the largest comes first. The next bit is set unless
+ * the member has a free extra partition, and the low 31 bits hold the member.
+ */
+ private long giverKey(int member) {
+ long notFree = freeExtraPartitions[member].isEmpty() ? 1 : 0;
+ return ((long) (Integer.MAX_VALUE - quotas.assignmentSize(member)) << 32) | notFree << 31 | member;
+ }
+
+ /**
+ * Returns the key of the member in the min heaps of receivers. The high 32 bits hold its
+ * assignment size. The next bit is set unless the member owns more than the minimum quota of a
+ * topic with extra partitions, and the low 31 bits hold the member.
+ */
+ private long receiverKey(int member) {
+ long notOwning = ownedAboveMinQuota[member].isEmpty() ? 1 : 0;
+ return ((long) quotas.assignmentSize(member) << 32) | notOwning << 31 | member;
+ }
+
+ /**
+ * @return The member of a key of {@link #giverKey} or {@link #receiverKey}.
+ */
+ private static int member(long key) {
+ return (int) (key & Integer.MAX_VALUE);
+ }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/steps/BalanceExtraPartitionsTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/steps/BalanceExtraPartitionsTest.java
new file mode 100644
index 00000000000..d105586fba7
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/steps/BalanceExtraPartitionsTest.java
@@ -0,0 +1,689 @@
+/*
+ * 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.steps;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.coordinator.group.assignor.uniform2.GroupSpecFixture;
+import org.apache.kafka.coordinator.group.assignor.uniform2.QuotasFixture;
+import org.apache.kafka.coordinator.group.assignor.uniform2.TopicsFixture;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+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.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Tests of {@link BalanceExtraPartitions}.
+ */
+public class BalanceExtraPartitionsTest {
+ 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 T6 = new Uuid(6L, 1L);
+
+ /**
+ * Nothing may move in an empty group, whose single cohort has no members.
+ */
+ @Test
+ public void testNothingMayMoveInAnEmptyGroup() {
+ var fixture = new QuotasFixture(new GroupSpecFixture().build(), new TopicsFixture().withTopic(T1, 3).build());
+ assertFalse(BalanceExtraPartitions.mayMove(fixture.group, fixture.quotas));
+ }
+
+ /**
+ * A and B subscribe to T1 and T2, of 3 partitions each, so each topic has 1 extra partition. A
+ * has the one of T1 and B the one of T2. The assignment sizes are 3 and 3, so nothing may move.
+ */
+ @Test
+ public void testNothingMayMoveWhenAssignmentSizesAreWithinOne() {
+ var fixture = homogeneousGroup();
+ fixture.addExtras("A", T1);
+ fixture.addExtras("B", T2);
+ assertEquals(List.of(3, 3), fixture.assignmentSizes());
+
+ assertFalse(BalanceExtraPartitions.mayMove(fixture.group, fixture.quotas));
+ }
+
+ /**
+ * A subscribes to T1 and T2, B to T1 only, in two cohorts. T1, 3 partitions, has 1 extra
+ * partition, A's, and T2 gives A 1 more partition. The assignment sizes are 3 and 1, two apart,
+ * so something may move.
+ */
+ @Test
+ public void testMayMoveWhenAssignmentSizesAreTwoApart() {
+ var fixture = twoCohortGroup();
+ fixture.addExtras("A", T1);
+ assertEquals(List.of(3, 1), fixture.assignmentSizes());
+
+ assertTrue(BalanceExtraPartitions.mayMove(fixture.group, fixture.quotas));
+ }
+
+ /**
+ * A subscribes to T1 and T2, and B to T3 and T4, in two cohorts. The counts per cohort take 2
+ * members times 2 cohorts, 4 entries, as many as the 4 topics of the subscriptions, so they
+ * fit. When B subscribes to T3 only, the subscriptions have 3 topics, and the counts do not.
+ */
+ @Test
+ public void testCountsFitWhenTheyTakeNoMoreEntriesThanTheSubscriptions() {
+ var describer = new TopicsFixture()
+ .withTopic(T1, 1)
+ .withTopic(T2, 1)
+ .withTopic(T3, 1)
+ .withTopic(T4, 1)
+ .build();
+ var fitting = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T3, T4))
+ .build();
+ var notFitting = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T3))
+ .build();
+
+ assertTrue(BalanceExtraPartitions.countsFit(new QuotasFixture(fitting, describer).group));
+ assertFalse(BalanceExtraPartitions.countsFit(new QuotasFixture(notFitting, describer).group));
+ }
+
+ /**
+ * 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, three of which subscribe to each topic. T1, T2 and T3, of 5 partitions each, have a
+ * minimum quota of 1 and 2 extra partitions each. A has one of every topic, B one of T1 and one
+ * of T2, and C the second 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.
+ */
+ @Test
+ public void testBalanceMovesUntilNoExtraPartitionCanMove() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2, T3))
+ .withMember("B", Set.of(T1, T2))
+ .withMember("C", Set.of(T2, T3))
+ .withMember("D", Set.of(T1, T3))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 5)
+ .withTopic(T2, 5)
+ .withTopic(T3, 5)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1, T2, T3);
+ fixture.addExtras("B", T1, T2);
+ fixture.addExtras("C", T3);
+ assertEquals(List.of(6, 4, 3, 2), fixture.assignmentSizes());
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("A", "B"), T2, List.of("B", "C"), T3, List.of("C", "D")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(4, 4, 4, 3), fixture.assignmentSizes());
+ }
+
+ /**
+ * A subscribes to T1 and T2, B to T1 only, and C to T2 and T3. T1 has no extra partition, T2
+ * has 1, A's. The assignment sizes are 3, 1 and 2. A and B are two apart, so something may
+ * move, but B cannot take the extra partition of T2, and C is only one partition smaller, so
+ * nothing moves.
+ */
+ @Test
+ public void testNothingMovesWhenNoExtraPartitionCanMove() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T1))
+ .withMember("C", Set.of(T2, T3))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 2)
+ .withTopic(T2, 3)
+ .withTopic(T3, 1)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T2);
+ assertTrue(BalanceExtraPartitions.mayMove(fixture.group, fixture.quotas));
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T2, List.of("A")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(3, 1, 2), fixture.assignmentSizes());
+ }
+
+ /**
+ * A and B subscribe to T1, T2 and T3, and C to T1 and T4, in cohorts 0 and 1. T1, 5 partitions,
+ * has a minimum quota of 1 and 2 extra partitions, A's and B's. T2 and T3, 3 partitions over A
+ * and B, have 1 extra partition each, both A's. T4 gives C 2 more partitions. The assignment
+ * sizes are 6, 4 and 3. A gives the extra partition of T1 to C, the smallest, and the
+ * assignment sizes end at 5, 4 and 4. That takes one move, where balancing the cohort of A and
+ * B first would move an extra partition from A to B, and then one of T1 to C.
+ */
+ @Test
+ public void testTheLargestMemberGivesToTheSmallestWhateverItsCohort() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2, T3))
+ .withMember("B", Set.of(T1, T2, T3))
+ .withMember("C", Set.of(T1, T4))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 5)
+ .withTopic(T2, 3)
+ .withTopic(T3, 3)
+ .withTopic(T4, 2)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1, T2, T3);
+ fixture.addExtras("B", T1);
+ assertEquals(List.of(6, 4, 3), fixture.assignmentSizes());
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("B", "C"), T2, List.of("A"), T3, List.of("A")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(5, 4, 4), fixture.assignmentSizes());
+ }
+
+ /**
+ * A and B subscribe to T1 and T2, of 5 partitions each. A holds T1 [0, 1, 2] and T2 [0, 1], and
+ * B holds T1 [3, 4] and T2 [2, 3, 4]. C joins with the same subscription. Each topic has a
+ * minimum quota of 1 and 2 extra partitions, all kept by A and B, which own more than 1
+ * partition of each. The assignment sizes are 4, 4 and 2. A and B are equally large and own all
+ * their extra partitions, so A, first in member order, gives one to C. Giving either of its
+ * extra partitions revokes a partition, and its topics are tried from the last, so the one of
+ * T2 goes. The assignment sizes end at 3, 4 and 3.
+ */
+ @Test
+ public void testWithinACohortTheLargestGivesToTheSmallest() {
+ var topics = Set.of(T1, T2);
+ var spec = new GroupSpecFixture()
+ .withMember("A", topics, mkAssignment(mkTopicAssignment(T1, 0, 1, 2), mkTopicAssignment(T2, 0, 1)))
+ .withMember("B", topics, mkAssignment(mkTopicAssignment(T1, 3, 4), mkTopicAssignment(T2, 2, 3, 4)))
+ .withMember("C", topics)
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 5)
+ .withTopic(T2, 5)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1, T2);
+ fixture.addExtras("B", T1, T2);
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("A", "B"), T2, List.of("B", "C")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(3, 4, 3), fixture.assignmentSizes());
+ }
+
+ /**
+ * A, B and C subscribe to T1, T2 and T3, of 4 partitions each. Each topic has a minimum quota
+ * of 1 and 1 extra partition. A has the three extra partitions and owns none of their
+ * partitions. B owns 2 partitions of T2 and C 2 of T1. The assignment sizes are 6, 3 and 3. A
+ * first gives B, the first of the smallest in member order, the extra partition of T2, which
+ * revokes one partition fewer, since B keeps its 2 partitions. A then gives C the one of T1,
+ * for the same reason, and keeps the one of T3.
+ */
+ @Test
+ public void testTheGiverGivesAnExtraPartitionTheReceiverOwnsFirst() {
+ var topics = Set.of(T1, T2, T3);
+ var spec = new GroupSpecFixture()
+ .withMember("A", topics)
+ .withMember("B", topics, mkAssignment(mkTopicAssignment(T2, 0, 1)))
+ .withMember("C", topics, mkAssignment(mkTopicAssignment(T1, 0, 1)))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 4)
+ .withTopic(T2, 4)
+ .withTopic(T3, 4)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1, T2, T3);
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("C"), T2, List.of("B"), T3, List.of("A")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(4, 4, 4), fixture.assignmentSizes());
+ }
+
+ /**
+ * A subscribes to T1, T2 and T3, B to T1 and T2. T1 and T2 have 3 partitions over A and B, so 1
+ * extra partition each, both A's. T3 gives A 1 more partition. A owns 2 partitions of T2, so
+ * its extra partition of T2 is owned and the one of T1 is free. The assignment sizes are 5 and
+ * 2. A gives B the extra partition of T1, which revokes no partition, and keeps the one of T2.
+ * The assignment sizes end at 4 and 3.
+ */
+ @Test
+ public void testTheGiverGivesAFreeExtraPartitionBeforeAnOwnedOne() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2, T3), mkAssignment(mkTopicAssignment(T2, 0, 1)))
+ .withMember("B", Set.of(T1, T2))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .withTopic(T2, 3)
+ .withTopic(T3, 1)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1, T2);
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("B"), T2, List.of("A")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(4, 3), fixture.assignmentSizes());
+ }
+
+ /**
+ * A, B and C subscribe to T1 to T4, of 4 partitions each. Each topic has a minimum quota of 1
+ * and 1 extra partition. A has the extra partitions of T1 and T2 and owns 2 partitions of each.
+ * B has those of T3 and T4 and owns 2 partitions of T3 only. The assignment sizes are 6, 6 and
+ * 4. A and B are equally large, but B has a free extra partition, the one of T4, whose move
+ * revokes no partition. B gives it to C, and A keeps both of its extra partitions.
+ */
+ @Test
+ public void testOnEqualSizesAGiverWithAFreeExtraPartitionGivesFirst() {
+ var topics = Set.of(T1, T2, T3, T4);
+ var spec = new GroupSpecFixture()
+ .withMember("A", topics, mkAssignment(mkTopicAssignment(T1, 0, 1), mkTopicAssignment(T2, 0, 1)))
+ .withMember("B", topics, mkAssignment(mkTopicAssignment(T3, 0, 1)))
+ .withMember("C", topics)
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 4)
+ .withTopic(T2, 4)
+ .withTopic(T3, 4)
+ .withTopic(T4, 4)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1, T2);
+ fixture.addExtras("B", T3, T4);
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("A"), T2, List.of("A"), T3, List.of("B"), T4, List.of("C")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(6, 5, 5), fixture.assignmentSizes());
+ }
+
+ /**
+ * A subscribes to T1 and T2, B to T1 and T4, and C to T1 and T3, in cohorts 0, 1 and 2. T1, 4
+ * partitions over the three, has 1 extra partition, A's. T2 gives A 2 more partitions, T4 gives
+ * B 1, and T3 gives C 1. The assignment sizes are 4, 2 and 2. B and C are equally small and own
+ * nothing, so B, first in member order, takes the extra partition of T1.
+ */
+ @Test
+ public void testReceiversOfEqualSizesGoInMemberOrder() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T1, T4))
+ .withMember("C", Set.of(T1, T3))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 4)
+ .withTopic(T2, 2)
+ .withTopic(T3, 1)
+ .withTopic(T4, 1)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1);
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("B")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(3, 3, 2), fixture.assignmentSizes());
+ }
+
+ /**
+ * A subscribes to T1 and T2, B and C to T1 only. T1, 4 partitions, has a minimum quota of 1 and
+ * 1 extra partition, A's. T2 gives A 1 more partition. The assignment sizes are 3, 1 and 1. C
+ * owns 2 partitions of T1, more than the minimum quota, and B owns none. So C, rather than B,
+ * receives the extra partition, and then keeps both of its partitions.
+ */
+ @Test
+ public void testOnEqualSizesAReceiverOwningMoreThanTheMinQuotaReceivesFirst() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T1))
+ .withMember("C", Set.of(T1), mkAssignment(mkTopicAssignment(T1, 1, 2)))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 4)
+ .withTopic(T2, 1)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1);
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("C")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(2, 1, 2), fixture.assignmentSizes());
+ }
+
+ /**
+ * A subscribes to T1 to T4, B and C to T1 to T5. T1 to T4, 4 partitions over 3 subscribers
+ * each, have 1 extra partition each, all A's. T5 gives B and C 1 more partition. The assignment
+ * sizes are 8, 5 and 5. A gives to B, the first of the smallest, then to C, skipping the entry
+ * of B with its former assignment size. The assignment sizes end at 6, 6 and 6. Looking at B
+ * with its former assignment size would have stopped at 7, 6 and 5.
+ */
+ @Test
+ public void testMembersWhoseSizeChangedAreLookedAtWithTheirNewSize() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2, T3, T4))
+ .withMember("B", Set.of(T1, T2, T3, T4, T5))
+ .withMember("C", Set.of(T1, T2, T3, T4, T5))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 4)
+ .withTopic(T2, 4)
+ .withTopic(T3, 4)
+ .withTopic(T4, 4)
+ .withTopic(T5, 2)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1, T2, T3, T4);
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("A"), T2, List.of("A"), T3, List.of("C"), T4, List.of("B")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(6, 6, 6), fixture.assignmentSizes());
+ }
+
+ /**
+ * A subscribes to T1 and T4, and B and C to T1, T2 and T3. T1, 5 partitions over the three, has
+ * 2 extra partitions, A's and B's. T2 and T3, 3 partitions over B and C, have 1 extra partition
+ * each, both C's. T4 gives A 5 more partitions. The assignment sizes are 7, 4 and 5. B, the
+ * smallest, already has an extra partition of T1, the only topic of A's extra partitions, so C,
+ * the next, takes it. Then C, at 6, gives B its extra partition of T3. The assignment sizes end
+ * at 6, 5 and 5.
+ */
+ @Test
+ public void testTheSmallestMemberWhichCannotTakeIsPassedOver() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T4))
+ .withMember("B", Set.of(T1, T2, T3))
+ .withMember("C", Set.of(T1, T2, T3))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 5)
+ .withTopic(T2, 3)
+ .withTopic(T3, 3)
+ .withTopic(T4, 5)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1);
+ fixture.addExtras("B", T1);
+ fixture.addExtras("C", T2, T3);
+ assertEquals(List.of(7, 4, 5), fixture.assignmentSizes());
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("B", "C"), T2, List.of("C"), T3, List.of("B")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(6, 5, 5), fixture.assignmentSizes());
+ }
+
+ /**
+ * A subscribes to T1 and T3, B to T1 and T2, and C to T2 only. T1, 3 partitions over A and B,
+ * has 1 extra partition, A's. T2, 3 partitions over B and C, has 1, B's. T3 gives A 2 more
+ * 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 the extra partition of T1. The
+ * assignment sizes end at 3, 3 and 2. Without waking A, they would end at 4, 2 and 2, and the
+ * extra partition of T1 could still move from A to B.
+ */
+ @Test
+ public void testAWaitingGiverWakesUpWhenASubscriberOfItsTopicsGetsSmaller() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T3))
+ .withMember("B", Set.of(T1, T2))
+ .withMember("C", Set.of(T2))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .withTopic(T2, 3)
+ .withTopic(T3, 2)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1);
+ fixture.addExtras("B", T2);
+ assertEquals(List.of(4, 3, 1), fixture.assignmentSizes());
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("B"), T2, List.of("C")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(3, 3, 2), fixture.assignmentSizes());
+ }
+
+ /**
+ * A and B subscribe to T1, T2 and T3, and C and D to T4, T5 and T6, in two cohorts that share
+ * no topic and are balanced on their own. T1 to T3 have 5 partitions over A and B, so a minimum
+ * quota of 2 and 1 extra partition each, all A's. T4 to T6 have 3 partitions over C and D, so a
+ * minimum quota of 1 and 1 extra partition each, all C's. The assignment sizes are 9, 6, 6 and
+ * 3. A gives an extra partition to B, and at 8 has no receiver, since B is at 7. C gives one to
+ * D, and at 5 has no receiver either. The assignment sizes end at 8, 7, 5 and 4.
+ */
+ @Test
+ public void testCohortsSharingNoTopicAreBalancedOnTheirOwn() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2, T3))
+ .withMember("B", Set.of(T1, T2, T3))
+ .withMember("C", Set.of(T4, T5, T6))
+ .withMember("D", Set.of(T4, T5, T6))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 5)
+ .withTopic(T2, 5)
+ .withTopic(T3, 5)
+ .withTopic(T4, 3)
+ .withTopic(T5, 3)
+ .withTopic(T6, 3)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1, T2, T3);
+ fixture.addExtras("C", T4, T5, T6);
+ assertEquals(List.of(9, 6, 6, 3), fixture.assignmentSizes());
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(
+ Map.ofEntries(
+ Map.entry(T1, List.of("A")),
+ Map.entry(T2, List.of("A")),
+ Map.entry(T3, List.of("B")),
+ Map.entry(T4, List.of("C")),
+ Map.entry(T5, List.of("C")),
+ Map.entry(T6, List.of("D"))
+ ),
+ fixture.membersWithExtraPartition()
+ );
+ assertEquals(List.of(8, 7, 5, 4), fixture.assignmentSizes());
+ }
+
+ /**
+ * A and B subscribe to T1 to T4, of 3 partitions each: a minimum quota of 1 and 1 extra
+ * partition each, all A's. The assignment sizes are 8 and 4. A gives B two extra partitions in
+ * a row, the last of its topics first: T4, then T3. The assignment sizes end at 6 and 6.
+ */
+ @Test
+ public void testAGiverGivesSeveralExtraPartitionsInARowToTheSameReceiver() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2, T3, T4))
+ .withMember("B", Set.of(T1, T2, T3, T4))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .withTopic(T2, 3)
+ .withTopic(T3, 3)
+ .withTopic(T4, 3)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1, T2, T3, T4);
+ assertEquals(List.of(8, 4), fixture.assignmentSizes());
+
+ balanceAndAssertBalanced(fixture);
+
+ assertEquals(Map.of(T1, List.of("A"), T2, List.of("A"), T3, List.of("B"), T4, List.of("B")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(6, 6), fixture.assignmentSizes());
+ }
+
+ /**
+ * On random groups, keeping the counts per cohort of every member or of one member at a time
+ * give exactly the same quotas, and these have the balance property. The groups have up to 40
+ * cohorts, so that topics become full in a cohort and stop being full as extra partitions move.
+ */
+ @Test
+ public void testBothWaysOfCountingGiveTheSameQuotas() {
+ var seeds = new Random(7);
+ for (int run = 0; run < 1000; run++) {
+ long seed = seeds.nextLong();
+ Map<Uuid, List<String>> first = null;
+ for (boolean keepCounts : new boolean[] {true, false}) {
+ var fixture = randomGroup(new Random(seed));
+ new BalanceExtraPartitions(fixture.group, fixture.currentAssignment, fixture.quotas, keepCounts).run();
+ assertBalanced(fixture);
+ var extras = fixture.membersWithExtraPartition();
+ if (first == null) {
+ first = extras;
+ } else {
+ assertEquals(first, extras, "seed " + seed + ", keepCounts " + keepCounts);
+ }
+ }
+ }
+ }
+
+ /**
+ * @return A group of 10 to 40 members in 1 to 40 subscriptions of 1 to 10 topics of 1 to 40
+ * partitions, owning random partitions, with the extra partitions of every topic on
+ * random distinct subscribers.
+ */
+ private static QuotasFixture randomGroup(Random random) {
+ int topicCount = 1 + random.nextInt(10);
+ var topicIds = new ArrayList<Uuid>();
+ var partitionCounts = new HashMap<Uuid, Integer>();
+ for (int topic = 0; topic < topicCount; topic++) {
+ topicIds.add(new Uuid(topic + 1, 1L));
+ partitionCounts.put(topicIds.get(topic), 1 + random.nextInt(40));
+ }
+ var subscriptions = new ArrayList<Set<Uuid>>();
+ for (int i = 1 + random.nextInt(40); i > 0; i--) {
+ var topics = new HashSet<Uuid>();
+ for (var topicId : topicIds) {
+ if (random.nextInt(2) == 0) {
+ topics.add(topicId);
+ }
+ }
+ topics.add(topicIds.get(random.nextInt(topicCount)));
+ subscriptions.add(topics);
+ }
+ var memberSubscriptions = new TreeMap<String, Set<Uuid>>();
+ for (int i = 10 + random.nextInt(31); i > 0; i--) {
+ memberSubscriptions.put(String.format("M%02d", i), subscriptions.get(random.nextInt(subscriptions.size())));
+ }
+ var owned = new HashMap<String, Map<Uuid, Set<Integer>>>();
+ for (var topicId : topicIds) {
+ var subscribers = memberSubscriptions.keySet().stream().filter(m -> memberSubscriptions.get(m).contains(topicId)).toList();
+ for (int partition = 0; partition < partitionCounts.get(topicId) && !subscribers.isEmpty(); partition++) {
+ if (random.nextInt(5) < 3) {
+ var owner = subscribers.get(random.nextInt(subscribers.size()));
+ owned.computeIfAbsent(owner, m -> new HashMap<>()).computeIfAbsent(topicId, t -> new HashSet<>()).add(partition);
+ }
+ }
+ }
+ var spec = new GroupSpecFixture();
+ memberSubscriptions.forEach((memberId, topics) -> spec.withMember(memberId, topics, owned.getOrDefault(memberId, Map.of())));
+ var describer = new TopicsFixture();
+ partitionCounts.forEach(describer::withTopic);
+ var fixture = new QuotasFixture(spec.build(), describer.build());
+ for (int topic = 0; topic < fixture.group.topicCount(); topic++) {
+ var topicId = fixture.group.topicId(topic);
+ var subscribers = new ArrayList<>(memberSubscriptions.keySet().stream().filter(m -> memberSubscriptions.get(m).contains(topicId)).toList());
+ Collections.shuffle(subscribers, random);
+ for (int i = 0; i < fixture.quotas.extraPartitions(topic); i++) {
+ fixture.addExtras(subscribers.get(i), topicId);
+ }
+ }
+ return fixture;
+ }
+
+ /**
+ * Checks that every topic has its number of extra partitions on distinct subscribers, and that
+ * no member with one is at least two partitions larger than a subscriber without one.
+ */
+ private static void assertBalanced(QuotasFixture fixture) {
+ var quotas = fixture.quotas;
+ // The moves that the balance property rules out, as "topic: giver to receiver".
+ var possibleMoves = new ArrayList<String>();
+ for (int topic = 0; topic < fixture.group.topicCount(); topic++) {
+ var membersWithExtraPartition = new HashSet<Integer>();
+ var members = quotas.membersWithExtraPartition(topic);
+ for (int i = 0; i < members.size(); i++) {
+ assertTrue(fixture.group.subscribes(members.get(i), topic));
+ membersWithExtraPartition.add(members.get(i));
+ }
+ assertEquals(quotas.extraPartitions(topic), membersWithExtraPartition.size());
+ for (int giver : membersWithExtraPartition) {
+ for (int receiver = 0; receiver < fixture.group.memberCount(); receiver++) {
+ boolean canTake = fixture.group.subscribes(receiver, topic) && !membersWithExtraPartition.contains(receiver);
+ if (canTake && quotas.assignmentSize(giver) - quotas.assignmentSize(receiver) >= 2) {
+ possibleMoves.add(topic + ": " + fixture.group.memberId(giver) + " to " + fixture.group.memberId(receiver));
+ }
+ }
+ }
+ }
+ assertEquals(List.of(), possibleMoves, "assignment sizes " + fixture.assignmentSizes());
+ }
+
+ /**
+ * Runs the balance step on the quotas of the fixture and checks that they have the balance
+ * property.
+ */
+ private static void balanceAndAssertBalanced(QuotasFixture fixture) {
+ new BalanceExtraPartitions(fixture.group, fixture.currentAssignment, fixture.quotas).run();
+ assertBalanced(fixture);
+ }
+
+ private static QuotasFixture homogeneousGroup() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T1, T2))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .withTopic(T2, 3)
+ .build();
+ return new QuotasFixture(spec, describer);
+ }
+
+ private static QuotasFixture twoCohortGroup() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T1))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .withTopic(T2, 1)
+ .build();
+ return new QuotasFixture(spec, describer);
+ }
+}