dajac commented on code in PR #23732: URL: https://github.com/apache/kafka/pull/23732#discussion_r4204186782
########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/GroupModel.java: ########## @@ -0,0 +1,351 @@ +/* + * 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 org.apache.kafka.coordinator.group.api.assignor.SubscriptionType; +import org.apache.kafka.coordinator.group.assignor.uniform2.util.IntList; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * The members, the subscribed topics and the cohorts of a group, numbered by small integers. + * + * <p>Members are numbered in member id order and topics in topic id order, so that every + * decision taken on these numbers depends on the content of the input only, never on the + * iteration order of the collections of the spec. A cohort is a set of members having the same + * subscription; cohorts are numbered in the order of their first member. The topics are those + * subscribed by at least one member and which exist. + * + * <p>With a homogeneous subscription type, only the subscription of the first member is read, + * as the coordinator guarantees that every member has the same one, and the group is a single + * cohort. Otherwise every subscription is read once, and summarized by a hash of its topics + * that does not depend on their order, so that grouping the members into cohorts is linear in + * the size of the subscriptions. + */ +final class GroupModel { + private final String[] memberIds; + private final int[] cohortOfMember; + + private final Uuid[] topicIds; + private final Map<Uuid, Integer> topicIndices; + private final int[] partitionCounts; + private final int[] subscriberCounts; + private final int[][] cohortsOfTopic; + + private final int[][] cohortMembers; + private final int[][] cohortTopics; + + /** + * Reads the members and their subscriptions. + * + * @param groupSpec The group spec. + * @param describer The describer of the subscribed topics. + */ + GroupModel(GroupSpec groupSpec, SubscribedTopicDescriber describer) { + memberIds = groupSpec.memberIds().toArray(new String[0]); + Arrays.sort(memberIds); + cohortOfMember = new int[memberIds.length]; + + var encountered = new EncounteredTopics(describer); + List<int[]> cohorts; + if (groupSpec.subscriptionType() == SubscriptionType.HOMOGENEOUS || memberIds.length <= 1) { + var topics = new IntList(16); + if (memberIds.length > 0) { + encountered.read(groupSpec, memberIds[0], topics); + } + cohorts = List.of(topics.toArray()); + } else { + cohorts = readCohorts(groupSpec, encountered); + } + + topicIds = encountered.ids.toArray(new Uuid[0]); + Arrays.sort(topicIds); + topicIndices = new HashMap<>(2 * topicIds.length); + partitionCounts = new int[topicIds.length]; + var renumbered = new int[topicIds.length]; + for (int topic = 0; topic < topicIds.length; topic++) { + topicIndices.put(topicIds[topic], topic); + int encounteredIndex = encountered.indices.get(topicIds[topic]); + renumbered[encounteredIndex] = topic; + partitionCounts[topic] = encountered.partitionCounts.get(encounteredIndex); + } + + cohortMembers = membersByCohort(cohorts.size()); + cohortTopics = new int[cohorts.size()][]; + subscriberCounts = new int[topicIds.length]; + for (int cohort = 0; cohort < cohorts.size(); cohort++) { + var topics = cohorts.get(cohort); + for (int i = 0; i < topics.length; i++) { + topics[i] = renumbered[topics[i]]; + subscriberCounts[topics[i]] += cohortMembers[cohort].length; + } + Arrays.sort(topics); + cohortTopics[cohort] = topics; + } + cohortsOfTopic = cohortsByTopic(); + } + + /** + * Numbers the topics in the order they are met, and remembers their partition count. + */ + private static final class EncounteredTopics { Review Comment: We may be able to find a better name for this one. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/GroupModel.java: ########## @@ -0,0 +1,351 @@ +/* + * 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 org.apache.kafka.coordinator.group.api.assignor.SubscriptionType; +import org.apache.kafka.coordinator.group.assignor.uniform2.util.IntList; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * The members, the subscribed topics and the cohorts of a group, numbered by small integers. + * + * <p>Members are numbered in member id order and topics in topic id order, so that every + * decision taken on these numbers depends on the content of the input only, never on the + * iteration order of the collections of the spec. A cohort is a set of members having the same + * subscription; cohorts are numbered in the order of their first member. The topics are those + * subscribed by at least one member and which exist. + * + * <p>With a homogeneous subscription type, only the subscription of the first member is read, + * as the coordinator guarantees that every member has the same one, and the group is a single + * cohort. Otherwise every subscription is read once, and summarized by a hash of its topics + * that does not depend on their order, so that grouping the members into cohorts is linear in + * the size of the subscriptions. + */ +final class GroupModel { + private final String[] memberIds; + private final int[] cohortOfMember; + + private final Uuid[] topicIds; + private final Map<Uuid, Integer> topicIndices; + private final int[] partitionCounts; + private final int[] subscriberCounts; + private final int[][] cohortsOfTopic; + + private final int[][] cohortMembers; + private final int[][] cohortTopics; + + /** + * Reads the members and their subscriptions. + * + * @param groupSpec The group spec. + * @param describer The describer of the subscribed topics. + */ + GroupModel(GroupSpec groupSpec, SubscribedTopicDescriber describer) { + memberIds = groupSpec.memberIds().toArray(new String[0]); + Arrays.sort(memberIds); + cohortOfMember = new int[memberIds.length]; + + var encountered = new EncounteredTopics(describer); + List<int[]> cohorts; + if (groupSpec.subscriptionType() == SubscriptionType.HOMOGENEOUS || memberIds.length <= 1) { + var topics = new IntList(16); + if (memberIds.length > 0) { + encountered.read(groupSpec, memberIds[0], topics); + } + cohorts = List.of(topics.toArray()); + } else { + cohorts = readCohorts(groupSpec, encountered); + } + + topicIds = encountered.ids.toArray(new Uuid[0]); + Arrays.sort(topicIds); + topicIndices = new HashMap<>(2 * topicIds.length); + partitionCounts = new int[topicIds.length]; + var renumbered = new int[topicIds.length]; + for (int topic = 0; topic < topicIds.length; topic++) { + topicIndices.put(topicIds[topic], topic); + int encounteredIndex = encountered.indices.get(topicIds[topic]); + renumbered[encounteredIndex] = topic; + partitionCounts[topic] = encountered.partitionCounts.get(encounteredIndex); + } + + cohortMembers = membersByCohort(cohorts.size()); + cohortTopics = new int[cohorts.size()][]; + subscriberCounts = new int[topicIds.length]; + for (int cohort = 0; cohort < cohorts.size(); cohort++) { + var topics = cohorts.get(cohort); + for (int i = 0; i < topics.length; i++) { + topics[i] = renumbered[topics[i]]; + subscriberCounts[topics[i]] += cohortMembers[cohort].length; + } + Arrays.sort(topics); + cohortTopics[cohort] = topics; + } + cohortsOfTopic = cohortsByTopic(); + } + + /** + * Numbers the topics in the order they are met, and remembers their partition count. + */ + private static final class EncounteredTopics { Review Comment: We usually put internal classes on top if not mistaken. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/GroupModel.java: ########## @@ -0,0 +1,351 @@ +/* + * 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 org.apache.kafka.coordinator.group.api.assignor.SubscriptionType; +import org.apache.kafka.coordinator.group.assignor.uniform2.util.IntList; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * The members, the subscribed topics and the cohorts of a group, numbered by small integers. + * + * <p>Members are numbered in member id order and topics in topic id order, so that every + * decision taken on these numbers depends on the content of the input only, never on the + * iteration order of the collections of the spec. A cohort is a set of members having the same + * subscription; cohorts are numbered in the order of their first member. The topics are those + * subscribed by at least one member and which exist. + * + * <p>With a homogeneous subscription type, only the subscription of the first member is read, + * as the coordinator guarantees that every member has the same one, and the group is a single + * cohort. Otherwise every subscription is read once, and summarized by a hash of its topics + * that does not depend on their order, so that grouping the members into cohorts is linear in + * the size of the subscriptions. + */ +final class GroupModel { + private final String[] memberIds; + private final int[] cohortOfMember; + + private final Uuid[] topicIds; + private final Map<Uuid, Integer> topicIndices; + private final int[] partitionCounts; + private final int[] subscriberCounts; + private final int[][] cohortsOfTopic; + + private final int[][] cohortMembers; + private final int[][] cohortTopics; + + /** + * Reads the members and their subscriptions. + * + * @param groupSpec The group spec. + * @param describer The describer of the subscribed topics. + */ + GroupModel(GroupSpec groupSpec, SubscribedTopicDescriber describer) { + memberIds = groupSpec.memberIds().toArray(new String[0]); + Arrays.sort(memberIds); + cohortOfMember = new int[memberIds.length]; + + var encountered = new EncounteredTopics(describer); + List<int[]> cohorts; + if (groupSpec.subscriptionType() == SubscriptionType.HOMOGENEOUS || memberIds.length <= 1) { + var topics = new IntList(16); + if (memberIds.length > 0) { + encountered.read(groupSpec, memberIds[0], topics); + } + cohorts = List.of(topics.toArray()); + } else { + cohorts = readCohorts(groupSpec, encountered); + } + + topicIds = encountered.ids.toArray(new Uuid[0]); + Arrays.sort(topicIds); + topicIndices = new HashMap<>(2 * topicIds.length); Review Comment: Should we use our new helper to create the HashMap with the correct size here? ########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/AssignmentTestUtils.java: ########## @@ -0,0 +1,80 @@ +/* + * 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.SubscriptionType; +import org.apache.kafka.coordinator.group.modern.Assignment; +import org.apache.kafka.coordinator.group.modern.GroupSpecImpl; +import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** + * Shared helpers for the uniform2 tests: building groups. + */ +public final class AssignmentTestUtils { + private AssignmentTestUtils() { } + + /** + * @return A member without rack. + */ + public static MemberSubscriptionAndAssignmentImpl member(Set<Uuid> topics, Assignment assignment) { Review Comment: I am not a fan of all those helpers. I wonder whether we could replace them by builders. Let's check. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/GroupModel.java: ########## @@ -0,0 +1,351 @@ +/* + * 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 org.apache.kafka.coordinator.group.api.assignor.SubscriptionType; +import org.apache.kafka.coordinator.group.assignor.uniform2.util.IntList; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * The members, the subscribed topics and the cohorts of a group, numbered by small integers. + * + * <p>Members are numbered in member id order and topics in topic id order, so that every + * decision taken on these numbers depends on the content of the input only, never on the + * iteration order of the collections of the spec. A cohort is a set of members having the same + * subscription; cohorts are numbered in the order of their first member. The topics are those + * subscribed by at least one member and which exist. + * + * <p>With a homogeneous subscription type, only the subscription of the first member is read, + * as the coordinator guarantees that every member has the same one, and the group is a single + * cohort. Otherwise every subscription is read once, and summarized by a hash of its topics + * that does not depend on their order, so that grouping the members into cohorts is linear in + * the size of the subscriptions. + */ +final class GroupModel { + private final String[] memberIds; + private final int[] cohortOfMember; + + private final Uuid[] topicIds; + private final Map<Uuid, Integer> topicIndices; + private final int[] partitionCounts; + private final int[] subscriberCounts; + private final int[][] cohortsOfTopic; + + private final int[][] cohortMembers; + private final int[][] cohortTopics; + + /** + * Reads the members and their subscriptions. + * + * @param groupSpec The group spec. + * @param describer The describer of the subscribed topics. + */ + GroupModel(GroupSpec groupSpec, SubscribedTopicDescriber describer) { + memberIds = groupSpec.memberIds().toArray(new String[0]); + Arrays.sort(memberIds); + cohortOfMember = new int[memberIds.length]; + + var encountered = new EncounteredTopics(describer); + List<int[]> cohorts; + if (groupSpec.subscriptionType() == SubscriptionType.HOMOGENEOUS || memberIds.length <= 1) { + var topics = new IntList(16); + if (memberIds.length > 0) { + encountered.read(groupSpec, memberIds[0], topics); + } + cohorts = List.of(topics.toArray()); + } else { + cohorts = readCohorts(groupSpec, encountered); + } + + topicIds = encountered.ids.toArray(new Uuid[0]); + Arrays.sort(topicIds); + topicIndices = new HashMap<>(2 * topicIds.length); + partitionCounts = new int[topicIds.length]; + var renumbered = new int[topicIds.length]; + for (int topic = 0; topic < topicIds.length; topic++) { + topicIndices.put(topicIds[topic], topic); + int encounteredIndex = encountered.indices.get(topicIds[topic]); + renumbered[encounteredIndex] = topic; + partitionCounts[topic] = encountered.partitionCounts.get(encounteredIndex); + } + + cohortMembers = membersByCohort(cohorts.size()); + cohortTopics = new int[cohorts.size()][]; + subscriberCounts = new int[topicIds.length]; + for (int cohort = 0; cohort < cohorts.size(); cohort++) { + var topics = cohorts.get(cohort); + for (int i = 0; i < topics.length; i++) { + topics[i] = renumbered[topics[i]]; + subscriberCounts[topics[i]] += cohortMembers[cohort].length; + } + Arrays.sort(topics); + cohortTopics[cohort] = topics; + } + cohortsOfTopic = cohortsByTopic(); + } + + /** + * Numbers the topics in the order they are met, and remembers their partition count. + */ + private static final class EncounteredTopics { + private final SubscribedTopicDescriber describer; + private final Map<Uuid, Integer> indices = new HashMap<>(); + private final List<Uuid> ids = new ArrayList<>(); + private final IntList partitionCounts = new IntList(16); + + EncounteredTopics(SubscribedTopicDescriber describer) { + this.describer = describer; + } + + /** + * Reads the subscription of the member. + * + * @param topics Receives the numbers of its topics that exist. + */ + void read(GroupSpec groupSpec, String memberId, IntList topics) { Review Comment: Why not passing the groupSpec via the constructor like the describer? Let's review other methods too. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/CurrentAssignment.java: ########## @@ -0,0 +1,182 @@ +/* + * 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.assignor.uniform2.util.IntList; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * The current assignment of a group: the owner of every partition. + * + * <p>A member owns a partition when it holds it, subscribes to its topic, and the partition + * exists. Everything else the members hold is stale: the partitions of topics they do not + * subscribe to or which no longer exist, the partitions beyond the partition count of their topic, + * and the empty sets of partitions. Stale partitions are dropped silently, and the members holding + * some get a new map. As the coordinator guarantees, a partition is held by at most one member. + */ +final class CurrentAssignment { + private final List<Map<Uuid, Set<Integer>>> assignments; + private final boolean[] holdsStalePartitions; + private final int[][] owners; + + /** + * The members and topics, as {@code member << 32 | topic}, of the sets holding partitions + * which do not exist, beyond the partition count or negative. They are rare, hence a set + * rather than an array. + */ + private final Set<Long> setsBeyondPartitionCount = new HashSet<>(); + + /** + * Reads the current assignment of every member. + * + * @param groupSpec The group spec. + * @param group The members and topics of the group. + */ + CurrentAssignment(GroupSpec groupSpec, GroupModel group) { + int memberCount = group.memberCount(); + int topicCount = group.topicCount(); + assignments = new ArrayList<>(memberCount); + for (int member = 0; member < memberCount; member++) { + assignments.add(groupSpec.memberAssignment(group.memberId(member)).partitions()); + } + holdsStalePartitions = new boolean[memberCount]; + owners = new int[topicCount][]; + + // The members are read cohort by cohort, so that marking the topics of the cohort tells + // in constant time whether a member subscribes to a topic. + var subscribed = new int[topicCount]; + Arrays.fill(subscribed, -1); + for (int cohort = 0; cohort < group.cohortCount(); cohort++) { + for (int topic : group.topicsOf(cohort)) { + subscribed[topic] = cohort; + } + for (int member : group.membersOf(cohort)) { + read(group, member, cohort, subscribed); + } + } + } + + /** + * Reads the current assignment of the member. The maps are iterated with {@code forEach}, + * which walks the table of a {@code HashMap} directly and is measurably cheaper than its + * iterator on large groups. A {@code HashSet} has no such {@code forEach}: reading its + * partitions allocates an iterator either way. + */ + private void read(GroupModel group, int member, int cohort, int[] subscribed) { + assignments.get(member).forEach((topicId, partitions) -> { + int topic = group.topicIndex(topicId); + if (topic < 0 || subscribed[topic] != cohort) { + holdsStalePartitions[member] = true; + return; + } + int owned = readPartitions(group, topic, member, partitions); + if (owned < partitions.size()) { + setsBeyondPartitionCount.add(key(member, topic)); + holdsStalePartitions[member] = true; + } else if (partitions.isEmpty()) { + holdsStalePartitions[member] = true; + } + }); + } + + /** + * Records the member as the owner of its partitions of the topic that exist. + * + * @return The number of partitions it owns. + */ + private int readPartitions(GroupModel group, int topic, int member, Set<Integer> partitions) { + int partitionCount = group.partitionCount(topic); + var owned = new int[1]; Review Comment: nit: This does not look great. Should we use a regular for loop so we can avoid it? -- 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]
