dajac commented on code in PR #23735: URL: https://github.com/apache/kafka/pull/23735#discussion_r4291910768
########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/steps/BalanceExtraPartitions.java: ########## @@ -0,0 +1,1050 @@ +/* + * 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 a single subscription, 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. + * + * @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); + var receiverKeys = new long[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); Review Comment: Good point. The heap keeps the array as its storage, and every move adds the giver and the receiver again, so a heap that started full grew at the first move. The array is now twice the size of the cohort, so the first moves fit without a copy. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/steps/BalanceExtraPartitions.java: ########## @@ -0,0 +1,1050 @@ +/* + * 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 a single subscription, 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. + * + * @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) { Review Comment: Right, it is exact only when every member subscribes to the same topics. With several subscriptions, assignment sizes can stay two apart while nothing can move, and the step then builds its state and moves nothing. An exact check needs more state per topic, so I would rather measure it on heterogeneous groups first and do it in a follow-up if it pays off. I added a comment for now. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
