dajac commented on code in PR #23736: URL: https://github.com/apache/kafka/pull/23736#discussion_r4348224783
########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/PartitionPlacement.java: ########## @@ -0,0 +1,201 @@ +/* + * 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; + +/** + * Places the partitions of a topic on its subscribers, given their quotas, moving as few + * partitions as possible. + * + * <p>Every owner keeps its partitions up to its quota, the lowest partition ids first. The other + * partitions, those without owner and those beyond the quota of their owner, go to the subscribers + * below their quota, in member order, the lowest partition ids first. A partition thus moves only + * when its owner's quota is below what it owns. + */ +final class PartitionPlacement { + /** + * The members and topics of the group. + */ + private final IndexedGroup group; + + /** + * The current assignment, which tells the owners of every partition. + */ + private final CurrentAssignment currentAssignment; + + /** + * The quotas of the members. + */ + private final Quotas quotas; + + /** + * Per member, the number of partitions of the topic being placed that it got so far. Reset + * after every topic. + */ + private final int[] placed; + + /** + * Per member, whether it gets an extra partition of the topic being placed. Reset after every + * topic. + */ + private final boolean[] hasExtraPartitionOfTopic; + + /** + * The partitions of the topic being placed that can go to the members below their quota, which + * are those without an owner and those beyond the quota of their owner. + */ + private final IntList free = new IntList(16); + + /** + * The subscribers of the topic being placed below their quota, in member order. + */ + private final IntList receivers = new IntList(16); + + /** + * The members that got partitions of the topic being placed, to reset {@link #placed}. + */ + private final IntList placedMembers = new IntList(16); + + /** + * @param group The members and topics of the group. + * @param currentAssignment The current assignment. + * @param quotas The quotas of the members. + */ + PartitionPlacement(IndexedGroup group, CurrentAssignment currentAssignment, Quotas quotas) { + this.group = group; + this.currentAssignment = currentAssignment; + this.quotas = quotas; + this.placed = new int[group.memberCount()]; + this.hasExtraPartitionOfTopic = new boolean[group.memberCount()]; + } + + /** + * Places the partitions of the topic. + * + * @param topic The topic. + * @param assignment Receives, for every partition of the topic, the member getting it. Its + * length is at least the partition count of the topic. + * @return {@code true} if a partition may move. {@code false} guarantees that every partition + * stays with its owner, and the caller need not read the assignment. + */ + boolean place(int topic, int[] assignment) { + markExtras(topic, true); + int minQuota = quotas.minQuota(topic); + var owners = currentAssignment.owners(topic); + free.clear(); + for (int partition = 0; partition < group.partitionCount(topic); partition++) { + int owner = owners == null ? -1 : owners[partition]; + if (owner >= 0 && placed[owner] < quota(owner, minQuota)) { + assignment[partition] = owner; + addPlaced(owner); + } else { + free.add(partition); + } + } + // The quotas add up to the partition count, so when every partition stays with its owner, + // every member has its quota and nothing moves. + boolean moves = !free.isEmpty(); + if (moves) { + collectReceivers(topic, minQuota); + int next = 0; + for (int i = 0; i < receivers.size(); i++) { + int member = receivers.get(i); + for (int missing = quota(member, minQuota) - placed[member]; missing > 0; missing--) { + assignment[free.get(next++)] = member; + } + } + } + reset(topic); + return moves; + } + + /** + * Sets or clears the marks of the members getting an extra partition of the topic. + */ + private void markExtras(int topic, boolean mark) { + var members = quotas.membersWithExtraPartition(topic); + for (int i = 0; i < members.size(); i++) { + hasExtraPartitionOfTopic[members.get(i)] = mark; + } + } + + /** + * @return The quota of the member for the topic being placed. + */ + private int quota(int member, int minQuota) { + return hasExtraPartitionOfTopic[member] ? minQuota + 1 : minQuota; + } Review Comment: Right, but `Quotas` is the output of the steps and the input of the placement, so I'd rather not keep the marks of the topic being placed in it. Answering in constant time from `Quotas` itself would require a flag per member and topic, or a per-member index kept up to date by every step that changes the quotas. Here, the marks are the placement's own working state, built once per topic, so `quota` stays next to them. ########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/PartitionPlacementTest.java: ########## @@ -0,0 +1,200 @@ +/* + * 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.Arrays; +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.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Tests of {@link PartitionPlacement}, on quotas set by hand. A topic with {@code P} partitions and + * {@code N} subscribers has a minimum quota of {@code P / N}, and each test gives its {@code P % N} + * extra partitions to chosen subscribers. + * + * <p>The member ids sort as A < B < C, so that A is member 0, B is member 1 and C is member + * 2, and the topic ids as T1 < T2, so that T1 is topic 0 and T2 is topic 1. + */ +public class PartitionPlacementTest { + private static final Uuid T1 = new Uuid(1L, 1L); + private static final Uuid T2 = new Uuid(2L, 1L); + + private static final int A = 0; + private static final int B = 1; + private static final int C = 2; + + /** + * Fills the assignment before placing, to tell whether {@code place} wrote it. + */ + private static final int UNTOUCHED = -7; Review Comment: `-1` works. The tests that I added in this round use it as well so I kept the constant. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/PartitionPlacement.java: ########## @@ -0,0 +1,201 @@ +/* + * 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; + +/** + * Places the partitions of a topic on its subscribers, given their quotas, moving as few + * partitions as possible. + * + * <p>Every owner keeps its partitions up to its quota, the lowest partition ids first. The other + * partitions, those without owner and those beyond the quota of their owner, go to the subscribers + * below their quota, in member order, the lowest partition ids first. A partition thus moves only + * when its owner's quota is below what it owns. + */ +final class PartitionPlacement { + /** + * The members and topics of the group. + */ + private final IndexedGroup group; + + /** + * The current assignment, which tells the owners of every partition. + */ + private final CurrentAssignment currentAssignment; + + /** + * The quotas of the members. + */ + private final Quotas quotas; + + /** + * Per member, the number of partitions of the topic being placed that it got so far. Reset + * after every topic. + */ + private final int[] placed; + + /** + * Per member, whether it gets an extra partition of the topic being placed. Reset after every + * topic. + */ + private final boolean[] hasExtraPartitionOfTopic; + + /** + * The partitions of the topic being placed that can go to the members below their quota, which + * are those without an owner and those beyond the quota of their owner. + */ + private final IntList free = new IntList(16); + + /** + * The subscribers of the topic being placed below their quota, in member order. + */ + private final IntList receivers = new IntList(16); + + /** + * The members that got partitions of the topic being placed, to reset {@link #placed}. + */ + private final IntList placedMembers = new IntList(16); + + /** + * @param group The members and topics of the group. + * @param currentAssignment The current assignment. + * @param quotas The quotas of the members. + */ + PartitionPlacement(IndexedGroup group, CurrentAssignment currentAssignment, Quotas quotas) { + this.group = group; + this.currentAssignment = currentAssignment; + this.quotas = quotas; + this.placed = new int[group.memberCount()]; + this.hasExtraPartitionOfTopic = new boolean[group.memberCount()]; + } + + /** + * Places the partitions of the topic. + * + * @param topic The topic. + * @param assignment Receives, for every partition of the topic, the member getting it. Its + * length is at least the partition count of the topic. + * @return {@code true} if a partition may move. {@code false} guarantees that every partition + * stays with its owner, and the caller need not read the assignment. + */ + boolean place(int topic, int[] assignment) { + markExtras(topic, true); + int minQuota = quotas.minQuota(topic); + var owners = currentAssignment.owners(topic); + free.clear(); + for (int partition = 0; partition < group.partitionCount(topic); partition++) { + int owner = owners == null ? -1 : owners[partition]; + if (owner >= 0 && placed[owner] < quota(owner, minQuota)) { + assignment[partition] = owner; + addPlaced(owner); + } else { + free.add(partition); + } + } + // The quotas add up to the partition count, so when every partition stays with its owner, + // every member has its quota and nothing moves. + boolean moves = !free.isEmpty(); + if (moves) { + collectReceivers(topic, minQuota); + int next = 0; + for (int i = 0; i < receivers.size(); i++) { + int member = receivers.get(i); + for (int missing = quota(member, minQuota) - placed[member]; missing > 0; missing--) { + assignment[free.get(next++)] = member; + } + } + } + reset(topic); + return moves; + } + + /** + * Sets or clears the marks of the members getting an extra partition of the topic. + */ + private void markExtras(int topic, boolean mark) { + var members = quotas.membersWithExtraPartition(topic); + for (int i = 0; i < members.size(); i++) { + hasExtraPartitionOfTopic[members.get(i)] = mark; + } + } + + /** + * @return The quota of the member for the topic being placed. + */ + private int quota(int member, int minQuota) { + return hasExtraPartitionOfTopic[member] ? minQuota + 1 : minQuota; + } + + /** + * Counts one more partition placed on the member. + */ + private void addPlaced(int member) { + if (placed[member]++ == 0) { + placedMembers.add(member); + } + } + + /** + * Collects in {@link #receivers} the subscribers of the topic below their quota, in member + * order. With a minimum quota of 0, only the members with an extra partition can be below their + * quota. + */ + private void collectReceivers(int topic, int minQuota) { + receivers.clear(); + if (minQuota == 0) { + var members = quotas.membersWithExtraPartition(topic); + for (int i = 0; i < members.size(); i++) { + if (placed[members.get(i)] == 0) { + receivers.add(members.get(i)); + } + } + } else { + for (int cohort : group.cohortsOf(topic)) { + for (int member : group.membersOf(cohort)) { + if (placed[member] < quota(member, minQuota)) { + receivers.add(member); + } + } + } + } + // The members of a single cohort are in member order already, unlike the members with an + // extra partition or the members of several cohorts. + if (minQuota == 0 || group.cohortsOf(topic).length > 1) { + receivers.sort(); + } Review Comment: Agreed, it is not worth it. The result stays deterministic without the sort: within a cohort the receivers come in member order, and with a min quota of 0 they come in the order of the members with an extra partition, which is deterministic too. Only which member gets which partition ids changes, not how many partitions each member gets. I removed it. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/PartitionPlacement.java: ########## @@ -0,0 +1,201 @@ +/* + * 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; + +/** + * Places the partitions of a topic on its subscribers, given their quotas, moving as few + * partitions as possible. + * + * <p>Every owner keeps its partitions up to its quota, the lowest partition ids first. The other + * partitions, those without owner and those beyond the quota of their owner, go to the subscribers + * below their quota, in member order, the lowest partition ids first. A partition thus moves only + * when its owner's quota is below what it owns. + */ +final class PartitionPlacement { + /** + * The members and topics of the group. + */ + private final IndexedGroup group; + + /** + * The current assignment, which tells the owners of every partition. + */ + private final CurrentAssignment currentAssignment; + + /** + * The quotas of the members. + */ + private final Quotas quotas; + + /** + * Per member, the number of partitions of the topic being placed that it got so far. Reset + * after every topic. + */ + private final int[] placed; + + /** + * Per member, whether it gets an extra partition of the topic being placed. Reset after every + * topic. + */ + private final boolean[] hasExtraPartitionOfTopic; + + /** + * The partitions of the topic being placed that can go to the members below their quota, which + * are those without an owner and those beyond the quota of their owner. + */ + private final IntList free = new IntList(16); + + /** + * The subscribers of the topic being placed below their quota, in member order. + */ + private final IntList receivers = new IntList(16); + + /** + * The members that got partitions of the topic being placed, to reset {@link #placed}. + */ + private final IntList placedMembers = new IntList(16); + + /** + * @param group The members and topics of the group. + * @param currentAssignment The current assignment. + * @param quotas The quotas of the members. + */ + PartitionPlacement(IndexedGroup group, CurrentAssignment currentAssignment, Quotas quotas) { + this.group = group; + this.currentAssignment = currentAssignment; + this.quotas = quotas; + this.placed = new int[group.memberCount()]; + this.hasExtraPartitionOfTopic = new boolean[group.memberCount()]; + } + + /** + * Places the partitions of the topic. + * + * @param topic The topic. + * @param assignment Receives, for every partition of the topic, the member getting it. Its + * length is at least the partition count of the topic. + * @return {@code true} if a partition may move. {@code false} guarantees that every partition + * stays with its owner, and the caller need not read the assignment. + */ + boolean place(int topic, int[] assignment) { + markExtras(topic, true); + int minQuota = quotas.minQuota(topic); + var owners = currentAssignment.owners(topic); + free.clear(); + for (int partition = 0; partition < group.partitionCount(topic); partition++) { + int owner = owners == null ? -1 : owners[partition]; + if (owner >= 0 && placed[owner] < quota(owner, minQuota)) { + assignment[partition] = owner; + addPlaced(owner); + } else { + free.add(partition); + } + } + // The quotas add up to the partition count, so when every partition stays with its owner, + // every member has its quota and nothing moves. + boolean moves = !free.isEmpty(); + if (moves) { + collectReceivers(topic, minQuota); Review Comment: Right, the list was only there to sort the receivers. Following your other comment, I dropped the member order, so we now hand out the partitions while iterating over the members, without collecting them first. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/PartitionPlacement.java: ########## @@ -0,0 +1,201 @@ +/* + * 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; + +/** + * Places the partitions of a topic on its subscribers, given their quotas, moving as few + * partitions as possible. + * + * <p>Every owner keeps its partitions up to its quota, the lowest partition ids first. The other + * partitions, those without owner and those beyond the quota of their owner, go to the subscribers + * below their quota, in member order, the lowest partition ids first. A partition thus moves only + * when its owner's quota is below what it owns. + */ +final class PartitionPlacement { + /** + * The members and topics of the group. + */ + private final IndexedGroup group; + + /** + * The current assignment, which tells the owners of every partition. + */ + private final CurrentAssignment currentAssignment; + + /** + * The quotas of the members. + */ + private final Quotas quotas; + + /** + * Per member, the number of partitions of the topic being placed that it got so far. Reset + * after every topic. + */ + private final int[] placed; + + /** + * Per member, whether it gets an extra partition of the topic being placed. Reset after every + * topic. + */ + private final boolean[] hasExtraPartitionOfTopic; Review Comment: Yeah, I agree. `Quotas` only keeps the list of members with an extra partition per topic. Answering this in constant time for any member and any topic would require a flag per member and per topic, so 50M flags with 10,000 members and 5,000 topics. This is why both steps mark the members of the topic that they work on. They are a bit different though: `AssignRemainingExtraPartitions` updates its marks while it hands out the extra partitions. I can extract the marking into a small helper next to `Quotas` in a follow-up. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/PartitionPlacement.java: ########## @@ -0,0 +1,201 @@ +/* + * 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; + +/** + * Places the partitions of a topic on its subscribers, given their quotas, moving as few + * partitions as possible. + * + * <p>Every owner keeps its partitions up to its quota, the lowest partition ids first. The other + * partitions, those without owner and those beyond the quota of their owner, go to the subscribers + * below their quota, in member order, the lowest partition ids first. A partition thus moves only + * when its owner's quota is below what it owns. + */ +final class PartitionPlacement { + /** + * The members and topics of the group. + */ + private final IndexedGroup group; + + /** + * The current assignment, which tells the owners of every partition. + */ + private final CurrentAssignment currentAssignment; + + /** + * The quotas of the members. + */ + private final Quotas quotas; + + /** + * Per member, the number of partitions of the topic being placed that it got so far. Reset + * after every topic. + */ + private final int[] placed; + + /** + * Per member, whether it gets an extra partition of the topic being placed. Reset after every + * topic. + */ + private final boolean[] hasExtraPartitionOfTopic; + + /** + * The partitions of the topic being placed that can go to the members below their quota, which + * are those without an owner and those beyond the quota of their owner. + */ + private final IntList free = new IntList(16); + + /** + * The subscribers of the topic being placed below their quota, in member order. + */ + private final IntList receivers = new IntList(16); + + /** + * The members that got partitions of the topic being placed, to reset {@link #placed}. + */ + private final IntList placedMembers = new IntList(16); + + /** + * @param group The members and topics of the group. + * @param currentAssignment The current assignment. + * @param quotas The quotas of the members. + */ + PartitionPlacement(IndexedGroup group, CurrentAssignment currentAssignment, Quotas quotas) { + this.group = group; + this.currentAssignment = currentAssignment; + this.quotas = quotas; + this.placed = new int[group.memberCount()]; + this.hasExtraPartitionOfTopic = new boolean[group.memberCount()]; + } + + /** + * Places the partitions of the topic. + * + * @param topic The topic. + * @param assignment Receives, for every partition of the topic, the member getting it. Its + * length is at least the partition count of the topic. + * @return {@code true} if a partition may move. {@code false} guarantees that every partition + * stays with its owner, and the caller need not read the assignment. + */ + boolean place(int topic, int[] assignment) { + markExtras(topic, true); + int minQuota = quotas.minQuota(topic); + var owners = currentAssignment.owners(topic); + free.clear(); + for (int partition = 0; partition < group.partitionCount(topic); partition++) { + int owner = owners == null ? -1 : owners[partition]; + if (owner >= 0 && placed[owner] < quota(owner, minQuota)) { + assignment[partition] = owner; + addPlaced(owner); + } else { + free.add(partition); + } + } + // The quotas add up to the partition count, so when every partition stays with its owner, + // every member has its quota and nothing moves. + boolean moves = !free.isEmpty(); + if (moves) { + collectReceivers(topic, minQuota); + int next = 0; + for (int i = 0; i < receivers.size(); i++) { + int member = receivers.get(i); + for (int missing = quota(member, minQuota) - placed[member]; missing > 0; missing--) { + assignment[free.get(next++)] = member; + } + } + } + reset(topic); + return moves; + } + + /** + * Sets or clears the marks of the members getting an extra partition of the topic. + */ + private void markExtras(int topic, boolean mark) { + var members = quotas.membersWithExtraPartition(topic); + for (int i = 0; i < members.size(); i++) { + hasExtraPartitionOfTopic[members.get(i)] = mark; + } + } + + /** + * @return The quota of the member for the topic being placed. + */ + private int quota(int member, int minQuota) { + return hasExtraPartitionOfTopic[member] ? minQuota + 1 : minQuota; + } + + /** + * Counts one more partition placed on the member. + */ + private void addPlaced(int member) { + if (placed[member]++ == 0) { + placedMembers.add(member); + } + } + + /** + * Collects in {@link #receivers} the subscribers of the topic below their quota, in member + * order. With a minimum quota of 0, only the members with an extra partition can be below their + * quota. + */ + private void collectReceivers(int topic, int minQuota) { + receivers.clear(); + if (minQuota == 0) { + var members = quotas.membersWithExtraPartition(topic); + for (int i = 0; i < members.size(); i++) { + if (placed[members.get(i)] == 0) { + receivers.add(members.get(i)); + } + } + } else { + for (int cohort : group.cohortsOf(topic)) { + for (int member : group.membersOf(cohort)) { + if (placed[member] < quota(member, minQuota)) { + receivers.add(member); + } + } + } + } + // The members of a single cohort are in member order already, unlike the members with an + // extra partition or the members of several cohorts. + if (minQuota == 0 || group.cohortsOf(topic).length > 1) { + receivers.sort(); + } + } + + /** + * Resets the numbers of partitions placed on the members, after a topic or before placing it + * again. + */ + private void clearPlaced() { + for (int i = 0; i < placedMembers.size(); i++) { + placed[placedMembers.get(i)] = 0; + } + placedMembers.clear(); + } + + /** + * Resets {@link #hasExtraPartitionOfTopic} and {@link #placed} after placing the topic. + */ + private void reset(int topic) { + markExtras(topic, false); + clearPlaced(); + } Review Comment: `kept` (previously `placed`) and `hasExtraPartitionOfTopic` have one entry per member, and `place` runs for every topic, including the stable ones. For instance, with 5,000 members and 500 topics of 40 partitions, a fill would clear 5,000 entries in each array for each of the 500 topics: 5M entries per assignment. Today, we only clear what we set: at most 80 entries per topic. For small groups, a fill would be as fast, but tracking the touched entries is never much slower. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/PartitionPlacement.java: ########## @@ -0,0 +1,201 @@ +/* + * 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; + +/** + * Places the partitions of a topic on its subscribers, given their quotas, moving as few + * partitions as possible. + * + * <p>Every owner keeps its partitions up to its quota, the lowest partition ids first. The other + * partitions, those without owner and those beyond the quota of their owner, go to the subscribers + * below their quota, in member order, the lowest partition ids first. A partition thus moves only + * when its owner's quota is below what it owns. + */ +final class PartitionPlacement { + /** + * The members and topics of the group. + */ + private final IndexedGroup group; + + /** + * The current assignment, which tells the owners of every partition. + */ + private final CurrentAssignment currentAssignment; + + /** + * The quotas of the members. + */ + private final Quotas quotas; + + /** + * Per member, the number of partitions of the topic being placed that it got so far. Reset + * after every topic. + */ + private final int[] placed; + + /** + * Per member, whether it gets an extra partition of the topic being placed. Reset after every + * topic. + */ + private final boolean[] hasExtraPartitionOfTopic; + + /** + * The partitions of the topic being placed that can go to the members below their quota, which + * are those without an owner and those beyond the quota of their owner. + */ + private final IntList free = new IntList(16); + + /** + * The subscribers of the topic being placed below their quota, in member order. + */ + private final IntList receivers = new IntList(16); + + /** + * The members that got partitions of the topic being placed, to reset {@link #placed}. + */ + private final IntList placedMembers = new IntList(16); + + /** + * @param group The members and topics of the group. + * @param currentAssignment The current assignment. + * @param quotas The quotas of the members. + */ + PartitionPlacement(IndexedGroup group, CurrentAssignment currentAssignment, Quotas quotas) { + this.group = group; + this.currentAssignment = currentAssignment; + this.quotas = quotas; + this.placed = new int[group.memberCount()]; + this.hasExtraPartitionOfTopic = new boolean[group.memberCount()]; + } + + /** + * Places the partitions of the topic. + * + * @param topic The topic. + * @param assignment Receives, for every partition of the topic, the member getting it. Its + * length is at least the partition count of the topic. + * @return {@code true} if a partition may move. {@code false} guarantees that every partition + * stays with its owner, and the caller need not read the assignment. + */ + boolean place(int topic, int[] assignment) { + markExtras(topic, true); + int minQuota = quotas.minQuota(topic); + var owners = currentAssignment.owners(topic); + free.clear(); + for (int partition = 0; partition < group.partitionCount(topic); partition++) { + int owner = owners == null ? -1 : owners[partition]; + if (owner >= 0 && placed[owner] < quota(owner, minQuota)) { + assignment[partition] = owner; + addPlaced(owner); + } else { + free.add(partition); + } + } + // The quotas add up to the partition count, so when every partition stays with its owner, + // every member has its quota and nothing moves. + boolean moves = !free.isEmpty(); + if (moves) { + collectReceivers(topic, minQuota); + int next = 0; + for (int i = 0; i < receivers.size(); i++) { + int member = receivers.get(i); + for (int missing = quota(member, minQuota) - placed[member]; missing > 0; missing--) { + assignment[free.get(next++)] = member; + } + } + } + reset(topic); + return moves; + } + + /** + * Sets or clears the marks of the members getting an extra partition of the topic. + */ + private void markExtras(int topic, boolean mark) { + var members = quotas.membersWithExtraPartition(topic); + for (int i = 0; i < members.size(); i++) { + hasExtraPartitionOfTopic[members.get(i)] = mark; + } + } + + /** + * @return The quota of the member for the topic being placed. + */ + private int quota(int member, int minQuota) { + return hasExtraPartitionOfTopic[member] ? minQuota + 1 : minQuota; + } + + /** + * Counts one more partition placed on the member. + */ + private void addPlaced(int member) { + if (placed[member]++ == 0) { + placedMembers.add(member); + } + } + + /** + * Collects in {@link #receivers} the subscribers of the topic below their quota, in member + * order. With a minimum quota of 0, only the members with an extra partition can be below their + * quota. + */ + private void collectReceivers(int topic, int minQuota) { + receivers.clear(); + if (minQuota == 0) { + var members = quotas.membersWithExtraPartition(topic); + for (int i = 0; i < members.size(); i++) { + if (placed[members.get(i)] == 0) { + receivers.add(members.get(i)); + } + } Review Comment: Yes, I think so. It happens as soon as a topic has fewer partitions than subscribers, so in large groups consuming many small topics. In our benchmark, it is the case for every topic of the homogeneous configurations with 5,000 members. Without it, we would walk all the subscribers of every topic that moves to find its members with an extra partition, which are at most as many as its partitions: with 5,000 members and topics of 40 partitions, 5,000 members instead of 40 per topic. -- 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]
