You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by da...@apache.org on 2026/10/15 07:46:33 UTC
(kafka) branch trunk updated: KAFKA-21214: [3/N] Add the quotas and the first two steps of the uniform2 assignor (#23733)
This is an automated email from the ASF dual-hosted git repository.
dajac 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 292ad36280a KAFKA-21214: [3/N] Add the quotas and the first two steps of the uniform2 assignor (#23733)
292ad36280a is described below
commit 292ad36280a55f4405033eb212451ea95513c6c1
Author: David Jacot <da...@apache.org>
AuthorDate: Thu Oct 15 09:46:26 2026 +0200
KAFKA-21214: [3/N] Add the quotas and the first two steps of the uniform2 assignor (#23733)
uniform2 decides how many partitions of every topic each member gets,
its quotas, in three steps, and then places the partitions. A topic with
P partitions and N subscribers has a minimum quota of P / N partitions,
and its P % N extra partitions go to as many distinct subscribers.
This patch adds the quotas and their first two steps:
- `Quotas` holds the minimum quota and the extra partitions of every
topic, the members getting the extra partitions, and the assignment
sizes of the members. The steps give the extra partitions and move
them.
- `KeepStickyExtraPartitions`: a member owning more than the minimum
quota of a topic keeps one of its extra partitions, the members with
the smallest assignments first when the owners outnumber the extra
partitions.
- `AssignRemainingExtraPartitions` gives the extra partitions that
nobody kept, one at a time, to subscribers without one of the topic,
among the ones with the smallest assignments: the topics with the
fewest subscribers first, and round robin within a cohort. It returns
without allocating anything when the owners kept them all.
- `IntPair` packs two ints in a long, so that pairs sort in a `long[]`
or a `LongHeap` without allocating.
- A step is created with what it works on and run once; the assignor
runs the steps in order in a later patch.
- `QuotasFixture` holds the indexed group, the current assignment and
the quotas of a group for the tests of the steps.
The balance step comes in a later patch.
Reviewers: Sean Quah <sq...@confluent.io>
---
.../uniform2/AssignRemainingExtraPartitions.java | 267 ++++++++++++++++++++
.../uniform2/KeepStickyExtraPartitions.java | 108 +++++++++
.../group/assignor/uniform2/Quotas.java | 169 +++++++++++++
.../group/assignor/uniform2/util/IntPair.java | 52 ++++
.../AssignRemainingExtraPartitionsTest.java | 269 +++++++++++++++++++++
.../uniform2/KeepStickyExtraPartitionsTest.java | 120 +++++++++
.../group/assignor/uniform2/QuotasFixture.java | 137 +++++++++++
.../group/assignor/uniform2/QuotasTest.java | 149 ++++++++++++
.../group/assignor/uniform2/util/IntPairTest.java | 85 +++++++
9 files changed, 1356 insertions(+)
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignRemainingExtraPartitions.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignRemainingExtraPartitions.java
new file mode 100644
index 00000000000..61b747c14cb
--- /dev/null
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignRemainingExtraPartitions.java
@@ -0,0 +1,267 @@
+/*
+ * 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.coordinator.group.assignor.uniform2.util.IntPair;
+import org.apache.kafka.coordinator.group.assignor.uniform2.util.LongHeap;
+
+import java.util.Arrays;
+
+/**
+ * The second step deciding the {@link Quotas}: assigns the extra partitions that no owner kept,
+ * one at a time, to subscribers not having one of the topic yet, among the ones with the smallest
+ * assignments.
+ *
+ * <p>The topics with the fewest subscribers go first, so that they can still choose the smallest
+ * of their few subscribers. When the subscriptions of the cohorts are subsets of each other, this
+ * leads to more balanced assignment sizes. The extra partition goes to the cohort holding the
+ * smallest assignment, and among its members those with the smallest assignment are served round
+ * robin in member order. The member with the smallest assignment may already have an extra
+ * partition of the topic, and another member of the cohort then gets it. The balance step moves it
+ * later if it has to. With an empty current assignment, a homogeneous group is served round robin.
+ *
+ * <p>On exit, every extra partition of every topic is assigned, each to a distinct subscriber.
+ */
+final class AssignRemainingExtraPartitions {
+ /**
+ * The state of a cohort while the extra partitions are assigned.
+ */
+ private static final class CohortState {
+ /**
+ * The members of the cohort.
+ */
+ final int[] members;
+
+ /**
+ * The position in the members of the next member to serve, round robin.
+ */
+ int cursor;
+
+ /**
+ * The lowest number of extra partitions among the members.
+ */
+ int lowestExtraPartitions;
+
+ /**
+ * How many members have the lowest number of extra partitions.
+ */
+ int membersWithLowestExtraPartitions;
+
+ /**
+ * How many members have an extra partition of the topic being assigned.
+ */
+ int membersWithExtraPartitionOfTopic;
+
+ CohortState(int[] members) {
+ this.members = members;
+ }
+ }
+
+ /**
+ * The members and topics of the group.
+ */
+ private final IndexedGroup group;
+
+ /**
+ * The quotas, which the step updates.
+ */
+ private final Quotas quotas;
+
+ /**
+ * The cohorts that can still receive an extra partition of the topic being assigned. They are
+ * ordered by the smallest assignment size among their members, whether that member has an
+ * extra partition of the topic or not, then by cohort order, see {@link #sizeAndCohort(int)}.
+ */
+ private LongHeap receivingCohorts;
+
+ /**
+ * Per cohort, its state.
+ */
+ private CohortState[] cohortStates;
+
+ /**
+ * Per member, whether it has an extra partition of the topic being assigned.
+ */
+ private boolean[] hasExtraPartitionOfTopic;
+
+ /**
+ * @param group The members and topics of the group.
+ * @param quotas The quotas, as the first step left them.
+ */
+ AssignRemainingExtraPartitions(IndexedGroup group, Quotas quotas) {
+ this.group = group;
+ this.quotas = quotas;
+ }
+
+ /**
+ * @return The key of the cohort in {@link #receivingCohorts}: the pair of the smallest
+ * assignment size among its members and the cohort, so that the cohorts are ordered by
+ * that assignment size, then cohort order. That assignment size does not change while
+ * the cohort is in the heap, since only the members of a polled cohort get an extra
+ * partition.
+ */
+ private long sizeAndCohort(int cohort) {
+ return IntPair.pack(quotas.minAssignmentSize(cohort) + cohortStates[cohort].lowestExtraPartitions, cohort);
+ }
+
+ /**
+ * Assigns the extra partitions that no owner kept. When the owners kept them all, as in a
+ * stable group, returns at once without allocating anything.
+ */
+ void run() {
+ if (!hasRemainingExtraPartitions()) {
+ return;
+ }
+ cohortStates = new CohortState[group.cohortCount()];
+ for (int cohort = 0; cohort < cohortStates.length; cohort++) {
+ cohortStates[cohort] = new CohortState(group.membersOf(cohort));
+ computeLowestExtraPartitions(cohortStates[cohort]);
+ }
+ hasExtraPartitionOfTopic = new boolean[group.memberCount()];
+ receivingCohorts = new LongHeap(16);
+ for (int topic : topicsWithRemainingExtraPartitions()) {
+ assignRemainingExtraPartitions(topic);
+ }
+ }
+
+ /**
+ * @return {@code true} if a topic has extra partitions that no owner kept.
+ */
+ private boolean hasRemainingExtraPartitions() {
+ for (int topic = 0; topic < group.topicCount(); topic++) {
+ if (quotas.membersWithExtraPartition(topic).size() < quotas.extraPartitions(topic)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /**
+ * Assigns the extra partitions of the topic that no owner kept, one at a time. Each goes to the
+ * next member of the cohort holding the smallest assignment, among the cohorts still having a
+ * member without one.
+ */
+ private void assignRemainingExtraPartitions(int topic) {
+ // The members with an extra partition of the topic. The list grows as they are assigned.
+ var members = quotas.membersWithExtraPartition(topic);
+ for (int i = 0; i < members.size(); i++) {
+ hasExtraPartitionOfTopic[members.get(i)] = true;
+ cohortStates[group.cohortOf(members.get(i))].membersWithExtraPartitionOfTopic++;
+ }
+ // The cohorts of the topic with a member still without an extra partition of it, the only
+ // ones that can receive one.
+ receivingCohorts.clear();
+ for (int cohort : group.cohortsOf(topic)) {
+ var state = cohortStates[cohort];
+ if (state.membersWithExtraPartitionOfTopic < state.members.length) {
+ receivingCohorts.add(sizeAndCohort(cohort));
+ }
+ }
+ // Computed before the loop, which adds to the members with an extra partition of the
+ // topic. There is at least one, since only the topics with some are assigned.
+ int remainingExtraPartitions = quotas.extraPartitions(topic) - members.size();
+ for (int i = 0; i < remainingExtraPartitions; i++) {
+ int cohort = IntPair.low(receivingCohorts.poll());
+ var state = cohortStates[cohort];
+ int member = nextMember(state);
+ int extraPartitionCount = quotas.extraPartitionCount(member);
+ quotas.giveExtraPartition(member, topic);
+ hasExtraPartitionOfTopic[member] = true;
+ // The last member of the cohort with the lowest number of extra partitions moves up.
+ if (extraPartitionCount == state.lowestExtraPartitions && --state.membersWithLowestExtraPartitions == 0) {
+ computeLowestExtraPartitions(state);
+ }
+ if (++state.membersWithExtraPartitionOfTopic < state.members.length) {
+ receivingCohorts.add(sizeAndCohort(cohort));
+ }
+ }
+ // Resets the state of the topic for the next one, including the members just assigned.
+ for (int i = 0; i < members.size(); i++) {
+ hasExtraPartitionOfTopic[members.get(i)] = false;
+ cohortStates[group.cohortOf(members.get(i))].membersWithExtraPartitionOfTopic = 0;
+ }
+ }
+
+ /**
+ * @return The topics with extra partitions that no owner kept, by increasing number of
+ * subscribers, then topic order.
+ */
+ private int[] topicsWithRemainingExtraPartitions() {
+ // Pairs of the number of subscribers and the topic, which sort by subscribers, then topic.
+ var subscribersAndTopics = new long[group.topicCount()];
+ int count = 0;
+ for (int topic = 0; topic < group.topicCount(); topic++) {
+ if (quotas.extraPartitions(topic) > quotas.membersWithExtraPartition(topic).size()) {
+ subscribersAndTopics[count++] = IntPair.pack(group.subscriberCount(topic), topic);
+ }
+ }
+ Arrays.sort(subscribersAndTopics, 0, count);
+ var topics = new int[count];
+ for (int i = 0; i < count; i++) {
+ topics[i] = IntPair.low(subscribersAndTopics[i]);
+ }
+ return topics;
+ }
+
+ /**
+ * Computes the lowest number of extra partitions of the members of the cohort, and how many
+ * of them have it.
+ */
+ private void computeLowestExtraPartitions(CohortState state) {
+ int lowest = Integer.MAX_VALUE;
+ int count = 0;
+ for (int member : state.members) {
+ int extraPartitionCount = quotas.extraPartitionCount(member);
+ if (extraPartitionCount < lowest) {
+ lowest = extraPartitionCount;
+ count = 1;
+ } else if (extraPartitionCount == lowest) {
+ count++;
+ }
+ }
+ state.lowestExtraPartitions = lowest;
+ state.membersWithLowestExtraPartitions = count;
+ }
+
+ /**
+ * @return The next member of the cohort, in round robin order from the cursor, having the
+ * lowest number of extra partitions among its members and no extra partition of the
+ * topic. If every member having the lowest number of extra partitions already has an
+ * extra partition of the topic, the first member without one having the fewest extra
+ * partitions.
+ */
+ private int nextMember(CohortState state) {
+ var members = state.members;
+ for (int step = 0; step < members.length; step++) {
+ int index = (state.cursor + step) % members.length;
+ int member = members[index];
+ if (quotas.extraPartitionCount(member) == state.lowestExtraPartitions && !hasExtraPartitionOfTopic[member]) {
+ state.cursor = (index + 1) % members.length;
+ return member;
+ }
+ }
+ // Every member with the lowest number of extra partitions already has one of the topic.
+ int best = -1;
+ for (int member : members) {
+ if (!hasExtraPartitionOfTopic[member] &&
+ (best < 0 || quotas.extraPartitionCount(member) < quotas.extraPartitionCount(best))) {
+ best = member;
+ }
+ }
+ return best;
+ }
+}
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/KeepStickyExtraPartitions.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/KeepStickyExtraPartitions.java
new file mode 100644
index 00000000000..5d3e7aa1e0c
--- /dev/null
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/KeepStickyExtraPartitions.java
@@ -0,0 +1,108 @@
+/*
+ * 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.coordinator.group.assignor.uniform2.util.IntList;
+import org.apache.kafka.coordinator.group.assignor.uniform2.util.IntPair;
+
+import java.util.Arrays;
+
+/**
+ * The first step deciding the {@link Quotas}: every member owning more than the minimum quota of a
+ * topic keeps one of its extra partitions, so that it only gives up what it owns beyond its quota.
+ * When more owners qualify than the topic has extra partitions, the ones with the smallest
+ * assignment sizes keep them. The assignment sizes include the extra partitions kept for the topics
+ * so far. Ties are broken by member order.
+ */
+final class KeepStickyExtraPartitions {
+ /**
+ * The members and topics of the group.
+ */
+ private final IndexedGroup group;
+
+ /**
+ * The current assignment, which tells the owners of every topic.
+ */
+ private final CurrentAssignment currentAssignment;
+
+ /**
+ * The quotas, which the step updates.
+ */
+ private final Quotas quotas;
+
+ /**
+ * @param group The members and topics of the group.
+ * @param currentAssignment The current assignment.
+ * @param quotas The quotas, at the minimum quota.
+ */
+ KeepStickyExtraPartitions(IndexedGroup group, CurrentAssignment currentAssignment, Quotas quotas) {
+ this.group = group;
+ this.currentAssignment = currentAssignment;
+ this.quotas = quotas;
+ }
+
+ /**
+ * Lets the owners of more than the minimum quota of every topic keep an extra partition.
+ */
+ void run() {
+ // Per member, the number of partitions of the topic that it owns.
+ var ownedPartitionCounts = new int[group.memberCount()];
+ // The owners of the topic.
+ var owners = new IntList(16);
+ // The owners above the minimum quota of the topic, which may keep an extra partition.
+ var ownersAboveMinQuota = new IntList(16);
+ for (int topic = 0; topic < group.topicCount(); topic++) {
+ int extraPartitions = quotas.extraPartitions(topic);
+ if (extraPartitions == 0) {
+ continue;
+ }
+ currentAssignment.countOwned(topic, ownedPartitionCounts, owners);
+ ownersAboveMinQuota.clear();
+ for (int i = 0; i < owners.size(); i++) {
+ int owner = owners.get(i);
+ if (ownedPartitionCounts[owner] > quotas.minQuota(topic)) {
+ ownersAboveMinQuota.add(owner);
+ }
+ // Reset for the next topic, as countOwned requires.
+ ownedPartitionCounts[owner] = 0;
+ }
+ // More owners qualify than the topic has extra partitions, as after a join. They are
+ // at most the subscribers of the topic, so the sort stays small.
+ if (ownersAboveMinQuota.size() > extraPartitions) {
+ sortBySize(ownersAboveMinQuota);
+ }
+ for (int i = 0; i < Math.min(ownersAboveMinQuota.size(), extraPartitions); i++) {
+ quotas.giveExtraPartition(ownersAboveMinQuota.get(i), topic);
+ }
+ }
+ }
+
+ /**
+ * Sorts the members by increasing assignment size, then member order.
+ */
+ private void sortBySize(IntList members) {
+ // Pairs of the assignment size and the member, which sort by assignment size, then member.
+ var assignmentSizesAndMembers = new long[members.size()];
+ for (int i = 0; i < assignmentSizesAndMembers.length; i++) {
+ assignmentSizesAndMembers[i] = IntPair.pack(quotas.assignmentSize(members.get(i)), members.get(i));
+ }
+ Arrays.sort(assignmentSizesAndMembers);
+ for (int i = 0; i < assignmentSizesAndMembers.length; i++) {
+ members.set(i, IntPair.low(assignmentSizesAndMembers[i]));
+ }
+ }
+}
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/Quotas.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/Quotas.java
new file mode 100644
index 00000000000..9afbcd0496a
--- /dev/null
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/Quotas.java
@@ -0,0 +1,169 @@
+/*
+ * 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.coordinator.group.assignor.uniform2.util.IntList;
+
+/**
+ * The quotas of the members: for every topic, how many of its partitions each subscriber gets.
+ *
+ * <p>A topic with {@code P} partitions and {@code N} subscribers has a minimum quota of
+ * {@code P / N} partitions and {@code P % N} extra partitions, which go to as many distinct
+ * subscribers, one each. The quota of a subscriber is thus the minimum quota, or one more when it
+ * gets an extra partition, and the assignment size of a member is the sum of its quotas.
+ *
+ * <p>The quotas start at the minimum quota. The steps deciding the quotas give the extra
+ * partitions with {@link #giveExtraPartition} and move them with {@link #moveExtraPartition},
+ * every extra partition going to a distinct subscriber of its topic.
+ */
+final class Quotas {
+ /**
+ * The empty list of members with an extra partition shared by all the topics without extra
+ * partitions. It never changes: {@link #giveExtraPartition} rejects such a topic, and
+ * {@link #moveExtraPartition} a giver without an extra partition of the topic.
+ */
+ private static final IntList NONE = new IntList(0);
+
+ /**
+ * The members and topics of the group.
+ */
+ private final IndexedGroup group;
+
+ /**
+ * Per topic, the minimum number of partitions that every subscriber gets.
+ */
+ private final int[] minQuotas;
+
+ /**
+ * Per topic, the number of subscribers getting one partition more than the minimum quota.
+ */
+ private final int[] extraPartitions;
+
+ /**
+ * Per cohort, the sum of the minimum quotas of its topics: the smallest assignment size of
+ * its members.
+ */
+ private final int[] minAssignmentSizes;
+
+ /**
+ * Per topic, the members getting one of its extra partitions.
+ */
+ private final IntList[] membersWithExtraPartition;
+
+ /**
+ * Per member, the number of topics of which it gets an extra partition.
+ */
+ private final int[] extraPartitionCounts;
+
+ /**
+ * Creates the quotas at the minimum quota: no member has an extra partition yet.
+ *
+ * @param group The members and topics of the group.
+ */
+ Quotas(IndexedGroup group) {
+ this.group = group;
+ int topicCount = group.topicCount();
+ minQuotas = new int[topicCount];
+ extraPartitions = new int[topicCount];
+ membersWithExtraPartition = new IntList[topicCount];
+ for (int topic = 0; topic < topicCount; topic++) {
+ minQuotas[topic] = group.partitionCount(topic) / group.subscriberCount(topic);
+ extraPartitions[topic] = group.partitionCount(topic) % group.subscriberCount(topic);
+ membersWithExtraPartition[topic] = extraPartitions[topic] == 0 ? NONE : new IntList(extraPartitions[topic]);
+ }
+ minAssignmentSizes = new int[group.cohortCount()];
+ for (int cohort = 0; cohort < group.cohortCount(); cohort++) {
+ for (int topic : group.topicsOf(cohort)) {
+ minAssignmentSizes[cohort] += minQuotas[topic];
+ }
+ }
+ extraPartitionCounts = new int[group.memberCount()];
+ }
+
+ /**
+ * @return The minimum number of partitions of the topic that every subscriber gets.
+ */
+ int minQuota(int topic) {
+ return minQuotas[topic];
+ }
+
+ /**
+ * @return The number of subscribers of the topic getting one partition more than the minimum
+ * quota.
+ */
+ int extraPartitions(int topic) {
+ return extraPartitions[topic];
+ }
+
+ /**
+ * @return The members getting an extra partition of the topic. The list must not be changed.
+ */
+ IntList membersWithExtraPartition(int topic) {
+ return membersWithExtraPartition[topic];
+ }
+
+ /**
+ * @return The number of topics of which the member gets an extra partition.
+ */
+ int extraPartitionCount(int member) {
+ return extraPartitionCounts[member];
+ }
+
+ /**
+ * @return The assignment size of the member: the number of partitions it gets over all its
+ * topics.
+ */
+ int assignmentSize(int member) {
+ return minAssignmentSizes[group.cohortOf(member)] + extraPartitionCounts[member];
+ }
+
+ /**
+ * @return The sum of the minimum quotas of the topics of the cohort: the smallest assignment
+ * size of its members.
+ */
+ int minAssignmentSize(int cohort) {
+ return minAssignmentSizes[cohort];
+ }
+
+ /**
+ * Gives an extra partition of the topic to the member. The member must not already have an
+ * extra partition of the topic.
+ *
+ * @throws IllegalStateException If every extra partition of the topic is given already.
+ */
+ void giveExtraPartition(int member, int topic) {
+ if (membersWithExtraPartition[topic].size() == extraPartitions[topic]) {
+ throw new IllegalStateException("No extra partition of topic " + group.topicId(topic) + " left for member "
+ + group.memberId(member));
+ }
+ membersWithExtraPartition[topic].add(member);
+ extraPartitionCounts[member]++;
+ }
+
+ /**
+ * Moves an extra partition of the topic from the giver to the receiver. The giver must have an
+ * extra partition of the topic and the receiver must not already have one.
+ *
+ * @throws IndexOutOfBoundsException If the giver has no extra partition of the topic.
+ */
+ void moveExtraPartition(int giver, int topic, int receiver) {
+ var members = membersWithExtraPartition[topic];
+ members.set(members.indexOf(giver), receiver);
+ extraPartitionCounts[giver]--;
+ extraPartitionCounts[receiver]++;
+ }
+}
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/util/IntPair.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/util/IntPair.java
new file mode 100644
index 00000000000..ee1fc0fc44a
--- /dev/null
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/util/IntPair.java
@@ -0,0 +1,52 @@
+/*
+ * 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.util;
+
+/**
+ * Packs a pair of ints in a long, without allocating: the high int in the high 32 bits of the long
+ * and the low int in the low 32 bits.
+ *
+ * <p>Pairs of non-negative ints sort as their longs do: by the high int, then the low int. The
+ * assignor sorts pairs this way in a {@code long[]} or a {@link LongHeap}, typically a key in the
+ * high int and the element it orders in the low int.
+ */
+public final class IntPair {
+ private IntPair() {}
+
+ /**
+ * @param high The high int.
+ * @param low The low int.
+ * @return The long holding the pair.
+ */
+ public static long pack(int high, int low) {
+ return ((long) high << 32) | (low & 0xFFFFFFFFL);
+ }
+
+ /**
+ * @return The high int of the pair.
+ */
+ public static int high(long pair) {
+ return (int) (pair >>> 32);
+ }
+
+ /**
+ * @return The low int of the pair.
+ */
+ public static int low(long pair) {
+ return (int) pair;
+ }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignRemainingExtraPartitionsTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignRemainingExtraPartitionsTest.java
new file mode 100644
index 00000000000..89868cca492
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignRemainingExtraPartitionsTest.java
@@ -0,0 +1,269 @@
+/*
+ * 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.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/**
+ * Tests of {@link AssignRemainingExtraPartitions}.
+ */
+public class AssignRemainingExtraPartitionsTest {
+ 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);
+
+ /**
+ * A subscribes to T1 and T2, B and D to T1 and T3, and C to T2 and T4. T1, 4 partitions over
+ * A, B and D, has 1 extra partition, and T2, 3 partitions over A and C, has 1 too; T3 and T4
+ * have none. The minimum quotas give assignment sizes of 2 to A, B and D, and 3 to C. T2 has
+ * the fewest subscribers and goes first: its extra partition goes to A, the smaller of A and
+ * C. Then the one of T1 goes to B, since A now has 3. The assignment sizes are within one of
+ * each other where the subscriptions allow it. Assigning T1 first would have given both extra
+ * partitions to A, the first cohort on ties: 4 partitions for A and 2 for B, which subscribes
+ * to T1.
+ */
+ @Test
+ public void testTopicsWithTheFewestSubscribersGoFirst() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T1, T3))
+ .withMember("C", Set.of(T2, T4))
+ .withMember("D", Set.of(T1, T3))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 4)
+ .withTopic(T2, 3)
+ .withTopic(T3, 2)
+ .withTopic(T4, 2)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+
+ new AssignRemainingExtraPartitions(fixture.group, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("B"), T2, List.of("A")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(3, 3, 3, 2), fixture.assignmentSizes());
+ }
+
+ /**
+ * A subscribes to T1 and T2, B to T1 and T3. T1, 3 partitions, has 1 extra partition. T2
+ * gives A 2 partitions and T3 gives B 1, so B has the smaller assignment and gets it.
+ */
+ @Test
+ public void testExtraPartitionGoesToTheCohortWithTheSmallestAssignment() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T1, T3))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .withTopic(T2, 2)
+ .withTopic(T3, 1)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+
+ new AssignRemainingExtraPartitions(fixture.group, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("B")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(3, 3), fixture.assignmentSizes());
+ }
+
+ /**
+ * A subscribes to T1 and T2, B to T1 and T3. T1, 3 partitions, has 1 extra partition. T2 and
+ * T3 have 1 partition each, so A and B have the same assignment size, and the extra partition
+ * of T1 goes to the first cohort, A's.
+ */
+ @Test
+ public void testTiesGoToTheFirstCohort() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T1, T3))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .withTopic(T2, 1)
+ .withTopic(T3, 1)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+
+ new AssignRemainingExtraPartitions(fixture.group, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("A")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(3, 2), fixture.assignmentSizes());
+ }
+
+ /**
+ * A, B and C subscribe to T1, T2 and T3, of 4 partitions each: 1 extra partition each. Each
+ * goes to the next member having no extra partition yet, and every member ends with one.
+ */
+ @Test
+ public void testMembersOfACohortAreServedInTurn() {
+ 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, 4)
+ .withTopic(T2, 4)
+ .withTopic(T3, 4)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+
+ new AssignRemainingExtraPartitions(fixture.group, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("A"), T2, List.of("B"), T3, List.of("C")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(4, 4, 4), fixture.assignmentSizes());
+ }
+
+ /**
+ * A, B and C subscribe to T1, of 4 partitions, and T2, of 5: 1 extra partition of T1 and 2
+ * of T2. A kept one of T2. The extra partition of T1 skips A, which has more extra partitions
+ * than B and C, and goes to B; the remaining one of T2 goes to C, the only member without
+ * any.
+ */
+ @Test
+ public void testMembersWithTheFewestExtraPartitionsAreServedFirst() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2))
+ .withMember("B", Set.of(T1, T2))
+ .withMember("C", Set.of(T1, T2))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 4)
+ .withTopic(T2, 5)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T2);
+
+ new AssignRemainingExtraPartitions(fixture.group, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("B"), T2, List.of("A", "C")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(3, 3, 3), fixture.assignmentSizes());
+ }
+
+ /**
+ * A, B and C subscribe to T1, T2 and T3, of 5 partitions each: 2 extra partitions each. A
+ * kept one of T1, and B and C both of T2 and T3. The last extra partition of T1 cannot go to
+ * A, the member with the fewest extra partitions, which has one of T1 already: it goes to the
+ * member without one having the fewest extra partitions, B before C in member order. The
+ * balance step evens the assignment sizes out later.
+ */
+ @Test
+ public void testAMemberWithMoreExtraPartitionsGetsOneWhenTheOthersHaveOneOfTheTopic() {
+ 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, 5)
+ .withTopic(T2, 5)
+ .withTopic(T3, 5)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1);
+ fixture.addExtras("B", T2, T3);
+ fixture.addExtras("C", T2, T3);
+
+ new AssignRemainingExtraPartitions(fixture.group, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("A", "B"), T2, List.of("B", "C"), T3, List.of("B", "C")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(4, 6, 5), fixture.assignmentSizes());
+ }
+
+ /**
+ * A subscribes to T1 only, B and C to T1 and T2. T1, 5 partitions, has 2 extra partitions,
+ * and A kept one. A has the smallest assignment, but its cohort has no member left without an
+ * extra partition of T1, and the other one goes to B.
+ */
+ @Test
+ public void testACohortWhoseMembersAllHaveAnExtraPartitionGetsNoMore() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1))
+ .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 fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1);
+
+ new AssignRemainingExtraPartitions(fixture.group, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("A", "B")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(2, 3, 2), fixture.assignmentSizes());
+ }
+
+ /**
+ * A and B subscribe to T1 only, C and D to T1 and T2. T1, 7 partitions over the four, has a
+ * minimum quota of 1 and 3 extra partitions, and T2 gives C and D 2 more partitions each. A and
+ * B have the smaller assignments and get the first two extra partitions. Their cohort is then
+ * full, and the third goes to C.
+ */
+ @Test
+ public void testACohortThatFillsUpGetsNoMore() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1))
+ .withMember("B", Set.of(T1))
+ .withMember("C", Set.of(T1, T2))
+ .withMember("D", Set.of(T1, T2))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 7)
+ .withTopic(T2, 4)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+
+ new AssignRemainingExtraPartitions(fixture.group, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("A", "B", "C")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(2, 2, 4, 3), fixture.assignmentSizes());
+ }
+
+ /**
+ * A and B subscribe to T1, of 3 partitions: a minimum quota of 1 and 1 extra partition, which
+ * A kept. Nothing is left to assign, and the quotas do not change.
+ */
+ @Test
+ public void testNothingLeftToAssign() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1))
+ .withMember("B", Set.of(T1))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+ fixture.addExtras("A", T1);
+
+ new AssignRemainingExtraPartitions(fixture.group, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("A")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(2, 1), fixture.assignmentSizes());
+ }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/KeepStickyExtraPartitionsTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/KeepStickyExtraPartitionsTest.java
new file mode 100644
index 00000000000..06575abbc80
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/KeepStickyExtraPartitionsTest.java
@@ -0,0 +1,120 @@
+/*
+ * 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.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+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;
+
+/**
+ * Tests of {@link KeepStickyExtraPartitions}.
+ */
+public class KeepStickyExtraPartitionsTest {
+ 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);
+
+ /**
+ * T1, 3 partitions over 3 members, has a minimum quota of 1 and no extra partition. T2, 8
+ * partitions, has a minimum quota of 2 and 2 extra partitions. A owns 3 partitions of T2 and C
+ * owns 4. Both own more than the minimum quota, so each keeps one extra partition, even C,
+ * which owns two partitions more than the minimum quota. The assignment sizes are 4, 3 and 4.
+ */
+ @Test
+ public void testOwnersAboveTheMinQuotaKeepAnExtraPartition() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2), mkAssignment(mkTopicAssignment(T1, 0), mkTopicAssignment(T2, 0, 1, 2)))
+ .withMember("B", Set.of(T1, T2), mkAssignment(mkTopicAssignment(T1, 1), mkTopicAssignment(T2, 3)))
+ .withMember("C", Set.of(T1, T2), mkAssignment(mkTopicAssignment(T1, 2), mkTopicAssignment(T2, 4, 5, 6, 7)))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .withTopic(T2, 8)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+
+ new KeepStickyExtraPartitions(fixture.group, fixture.currentAssignment, fixture.quotas).run();
+
+ assertEquals(Map.of(T2, List.of("A", "C")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(4, 3, 4), fixture.assignmentSizes());
+ }
+
+ /**
+ * T1 and T2 have 4 partitions each over A, B and C: a minimum quota of 1 and 1 extra
+ * partition each. A owns 2 partitions of both and B 2 of T1. T1 has two owners above the
+ * minimum quota for one extra partition, of the same assignment size so far: A keeps it,
+ * first in member order. A also keeps the one of T2, of which only A owns more than the
+ * minimum quota.
+ */
+ @Test
+ public void testOwnersOfTheSameSizeKeepInMemberOrder() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2), mkAssignment(mkTopicAssignment(T1, 0, 1), mkTopicAssignment(T2, 0, 1)))
+ .withMember("B", Set.of(T1, T2), mkAssignment(mkTopicAssignment(T1, 2, 3), mkTopicAssignment(T2, 2)))
+ .withMember("C", Set.of(T1, T2))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 4)
+ .withTopic(T2, 4)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+
+ new KeepStickyExtraPartitions(fixture.group, fixture.currentAssignment, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("A"), T2, List.of("A")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(4, 2, 2), fixture.assignmentSizes());
+ }
+
+ /**
+ * Three members own 3 of the 9 partitions of T1, T2 and T3 each; a fourth member joins. Each
+ * topic now has a minimum quota of 2 and 1 extra partition, and its 3 owners all own more than
+ * the minimum quota. The owners with the smallest assignments keep the extra partitions: A for
+ * T1; then B and C have the smaller assignment size for T2, B first in member order; then C
+ * for T3. D gets none, and the assignment sizes are within one of each other.
+ */
+ @Test
+ public void testOwnersWithTheSmallestAssignmentKeepTheExtraPartitions() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2, T3), mkAssignment(
+ mkTopicAssignment(T1, 0, 1, 2), mkTopicAssignment(T2, 0, 1, 2), mkTopicAssignment(T3, 0, 1, 2)))
+ .withMember("B", Set.of(T1, T2, T3), mkAssignment(
+ mkTopicAssignment(T1, 3, 4, 5), mkTopicAssignment(T2, 3, 4, 5), mkTopicAssignment(T3, 3, 4, 5)))
+ .withMember("C", Set.of(T1, T2, T3), mkAssignment(
+ mkTopicAssignment(T1, 6, 7, 8), mkTopicAssignment(T2, 6, 7, 8), mkTopicAssignment(T3, 6, 7, 8)))
+ .withMember("D", Set.of(T1, T2, T3))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 9)
+ .withTopic(T2, 9)
+ .withTopic(T3, 9)
+ .build();
+ var fixture = new QuotasFixture(spec, describer);
+
+ new KeepStickyExtraPartitions(fixture.group, fixture.currentAssignment, fixture.quotas).run();
+
+ assertEquals(Map.of(T1, List.of("A"), T2, List.of("B"), T3, List.of("C")), fixture.membersWithExtraPartition());
+ assertEquals(List.of(7, 7, 7, 6), fixture.assignmentSizes());
+ }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/QuotasFixture.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/QuotasFixture.java
new file mode 100644
index 00000000000..0f245b378a4
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/QuotasFixture.java
@@ -0,0 +1,137 @@
+/*
+ * 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.GroupSpec;
+import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.IntPredicate;
+
+/**
+ * The indexed group, the current assignment and the quotas of a group, for the tests of the steps
+ * deciding the quotas. The quotas start at the minimum quota, to which a test adds the extra
+ * partitions it starts from. The members and topics are referred to by their ids.
+ */
+final class QuotasFixture {
+ /**
+ * The members and topics of the group.
+ */
+ final IndexedGroup group;
+
+ /**
+ * The current assignment of the group.
+ */
+ final CurrentAssignment currentAssignment;
+
+ /**
+ * The quotas of the members.
+ */
+ final Quotas quotas;
+
+ /**
+ * Creates the fixture, with the quotas at the minimum quota.
+ *
+ * @param groupSpec The group spec.
+ * @param describer The describer of the topics.
+ */
+ QuotasFixture(GroupSpec groupSpec, SubscribedTopicDescriber describer) {
+ group = new IndexedGroup.Builder(groupSpec, describer).build();
+ currentAssignment = new CurrentAssignment(groupSpec, group);
+ quotas = new Quotas(group);
+ }
+
+ /**
+ * @return The index of the member.
+ */
+ int member(String memberId) {
+ for (int member = 0; member < group.memberCount(); member++) {
+ if (group.memberId(member).equals(memberId)) {
+ return member;
+ }
+ }
+ throw new IllegalArgumentException("Unknown member " + memberId);
+ }
+
+ /**
+ * @return The index of the topic.
+ */
+ int topic(Uuid topicId) {
+ return group.topicIndex(topicId);
+ }
+
+ /**
+ * Gives the member an extra partition of each of the topics. It must not have an extra
+ * partition of any of them yet.
+ */
+ void addExtras(String memberId, Uuid... topicIds) {
+ for (Uuid topicId : topicIds) {
+ if (quotas.membersWithExtraPartition(topic(topicId)).indexOf(member(memberId)) >= 0) {
+ throw new IllegalArgumentException(memberId + " already has an extra partition of " + topicId);
+ }
+ quotas.giveExtraPartition(member(memberId), topic(topicId));
+ }
+ }
+
+ /**
+ * @return Per topic, the members which have an extra partition of it, in member order. The
+ * topics where no member has an extra partition are left out.
+ */
+ Map<Uuid, List<String>> membersWithExtraPartition() {
+ var result = new HashMap<Uuid, List<String>>();
+ for (int topic = 0; topic < group.topicCount(); topic++) {
+ var members = quotas.membersWithExtraPartition(topic);
+ if (!members.isEmpty()) {
+ var memberIds = new ArrayList<String>();
+ for (int i = 0; i < members.size(); i++) {
+ memberIds.add(group.memberId(members.get(i)));
+ }
+ memberIds.sort(null);
+ result.put(group.topicId(topic), memberIds);
+ }
+ }
+ return result;
+ }
+
+ /**
+ * @return The members for which the predicate holds, in member id order.
+ */
+ List<String> members(IntPredicate predicate) {
+ var memberIds = new ArrayList<String>();
+ for (int member = 0; member < group.memberCount(); member++) {
+ if (predicate.test(member)) {
+ memberIds.add(group.memberId(member));
+ }
+ }
+ return memberIds;
+ }
+
+ /**
+ * @return The assignment sizes of the members, in member id order.
+ */
+ List<Integer> assignmentSizes() {
+ var sizes = new ArrayList<Integer>();
+ for (int member = 0; member < group.memberCount(); member++) {
+ sizes.add(quotas.assignmentSize(member));
+ }
+ return sizes;
+ }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/QuotasTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/QuotasTest.java
new file mode 100644
index 00000000000..41413a456b3
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/QuotasTest.java
@@ -0,0 +1,149 @@
+/*
+ * 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.assignor.uniform2.util.IntList;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Set;
+import java.util.function.IntUnaryOperator;
+import java.util.stream.IntStream;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+/**
+ * Tests of {@link Quotas}.
+ */
+public class QuotasTest {
+ 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 int A = 0;
+ private static final int B = 1;
+
+ /**
+ * 3 partitions over 2 subscribers give a minimum quota of 1 and 1 extra partition; 7 over 3
+ * a minimum quota of 2 and 1 extra partition; 2 over 3 a minimum quota of 0 and 2 extra
+ * partitions. The members of cohort 0, A and B, subscribe to all three topics and get at least
+ * 1 + 2 + 0 = 3 partitions each; the member of cohort 1, C, subscribes to T2 and T3 and gets
+ * at least 2 + 0 = 2.
+ */
+ @Test
+ public void testMinQuotaAndExtraPartitions() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1, T2, T3))
+ .withMember("B", Set.of(T1, T2, T3))
+ .withMember("C", Set.of(T2, T3))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .withTopic(T2, 7)
+ .withTopic(T3, 2)
+ .build();
+ var group = new IndexedGroup.Builder(spec, describer).build();
+
+ var quotas = new Quotas(group);
+
+ assertEquals(List.of(1, 2, 0), values(group.topicCount(), quotas::minQuota));
+ assertEquals(List.of(1, 1, 2), values(group.topicCount(), quotas::extraPartitions));
+ assertEquals(List.of(3, 2), values(group.cohortCount(), quotas::minAssignmentSize));
+ assertEquals(List.of(3, 3, 2), values(group.memberCount(), quotas::assignmentSize));
+ assertEquals(List.of(0, 0, 0), values(group.memberCount(), quotas::extraPartitionCount));
+ assertEquals(List.of(List.of(), List.of(), List.of()), membersWithExtraPartition(group, quotas));
+ }
+
+ /**
+ * A and B subscribe to T1, of 3 partitions, and T2, of 2. T1 has 1 extra partition, which A
+ * gets, and T2 has none. A second extra partition of T1 cannot be given, nor one of T2.
+ */
+ @Test
+ public void testGiveExtraPartition() {
+ 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, 2)
+ .build();
+ var group = new IndexedGroup.Builder(spec, describer).build();
+ var quotas = new Quotas(group);
+ int t1 = group.topicIndex(T1);
+ int t2 = group.topicIndex(T2);
+
+ quotas.giveExtraPartition(A, t1);
+
+ assertThrows(IllegalStateException.class, () -> quotas.giveExtraPartition(B, t1));
+ assertThrows(IllegalStateException.class, () -> quotas.giveExtraPartition(A, t2));
+ assertEquals(List.of(3, 2), values(group.memberCount(), quotas::assignmentSize));
+ assertEquals(List.of(1, 0), values(group.memberCount(), quotas::extraPartitionCount));
+ assertEquals(List.of(List.of(A), List.of()), membersWithExtraPartition(group, quotas));
+ }
+
+ /**
+ * A and B subscribe to T1, of 3 partitions, which has 1 extra partition, A's. B has none to
+ * give. Moving the extra partition from A to B gives B the larger assignment size.
+ */
+ @Test
+ public void testMoveExtraPartition() {
+ var spec = new GroupSpecFixture()
+ .withMember("A", Set.of(T1))
+ .withMember("B", Set.of(T1))
+ .build();
+ var describer = new TopicsFixture()
+ .withTopic(T1, 3)
+ .build();
+ var group = new IndexedGroup.Builder(spec, describer).build();
+ var quotas = new Quotas(group);
+ int t1 = group.topicIndex(T1);
+ quotas.giveExtraPartition(A, t1);
+
+ assertThrows(IndexOutOfBoundsException.class, () -> quotas.moveExtraPartition(B, t1, A));
+ quotas.moveExtraPartition(A, t1, B);
+
+ assertEquals(List.of(1, 2), values(group.memberCount(), quotas::assignmentSize));
+ assertEquals(List.of(0, 1), values(group.memberCount(), quotas::extraPartitionCount));
+ assertEquals(List.of(List.of(B)), membersWithExtraPartition(group, quotas));
+ }
+
+ /**
+ * @return The values of the numbers from 0 to the count, in order.
+ */
+ private static List<Integer> values(int count, IntUnaryOperator value) {
+ return IntStream.range(0, count).map(value).boxed().toList();
+ }
+
+ /**
+ * @return Per topic, the members getting one of its extra partitions, in the order they got
+ * them.
+ */
+ private static List<List<Integer>> membersWithExtraPartition(IndexedGroup group, Quotas quotas) {
+ return IntStream.range(0, group.topicCount())
+ .mapToObj(topic -> toList(quotas.membersWithExtraPartition(topic)))
+ .toList();
+ }
+
+ private static List<Integer> toList(IntList list) {
+ return Arrays.stream(list.toArray()).boxed().toList();
+ }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/util/IntPairTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/util/IntPairTest.java
new file mode 100644
index 00000000000..6ab88e97c03
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/util/IntPairTest.java
@@ -0,0 +1,85 @@
+/*
+ * 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.util;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/**
+ * Tests the pairs of ints packed in a long.
+ */
+public class IntPairTest {
+ /**
+ * Pairs covering zero, negative ints and the extremes of both ints.
+ */
+ private static final List<int[]> PAIRS = List.of(
+ new int[] {3, 7},
+ new int[] {0, 0},
+ new int[] {0, -1},
+ new int[] {-1, 0},
+ new int[] {Integer.MAX_VALUE, Integer.MIN_VALUE},
+ new int[] {Integer.MIN_VALUE, Integer.MAX_VALUE}
+ );
+
+ @Test
+ public void testPack() {
+ assertEquals((3L << 32) | 7L, IntPair.pack(3, 7));
+ assertEquals(0xFFFFFFFFL, IntPair.pack(0, -1));
+ assertEquals(-1L, IntPair.pack(-1, -1));
+ }
+
+ @Test
+ public void testHigh() {
+ for (int[] pair : PAIRS) {
+ assertEquals(pair[0], IntPair.high(IntPair.pack(pair[0], pair[1])));
+ }
+ }
+
+ @Test
+ public void testLow() {
+ for (int[] pair : PAIRS) {
+ assertEquals(pair[1], IntPair.low(IntPair.pack(pair[0], pair[1])));
+ }
+ }
+
+ /**
+ * Pairs of non-negative ints sort by the high int, then the low int.
+ */
+ @Test
+ public void testPairsOfNonNegativeIntsSortByHighThenLow() {
+ long[] pairs = {
+ IntPair.pack(2, 0),
+ IntPair.pack(1, 5),
+ IntPair.pack(1, 2),
+ IntPair.pack(0, Integer.MAX_VALUE)
+ };
+
+ Arrays.sort(pairs);
+
+ assertArrayEquals(new long[] {
+ IntPair.pack(0, Integer.MAX_VALUE),
+ IntPair.pack(1, 2),
+ IntPair.pack(1, 5),
+ IntPair.pack(2, 0)
+ }, pairs);
+ }
+}