squah-confluent commented on code in PR #23732: URL: https://github.com/apache/kafka/pull/23732#discussion_r4249274928
########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/CurrentAssignment.java: ########## @@ -0,0 +1,162 @@ +/* + * 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.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 and subscribes to its topic. The partitions of + * topics that a member does not subscribe to, or which no longer exist, are stale: they 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, a topic never has fewer partitions than its members hold, and no + * set of partitions is empty. + */ +final class CurrentAssignment { + /** + * Per member, the partitions it holds, as given by the spec. + */ + private final List<Map<Uuid, Set<Integer>>> assignments; + + /** + * Per member, whether it needs a new map because it holds stale partitions. + */ + private final boolean[] holdsStalePartitions; + + /** + * Per topic, the owner of every partition, or -1; null while no partition of the topic has an + * owner, so that no array is allocated for a topic that nobody owns, as on the first assignment + * of a group. + */ + private final int[][] owners; + + /** + * Reads the current assignment of every member. + * + * @param groupSpec The group spec. + * @param group The members and topics of the group. + */ + CurrentAssignment(GroupSpec groupSpec, IndexedGroup 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)) { + recordAssignment(group, member, cohort, subscribed); + } + } + } + + /** + * Records the owners of the partitions that the member holds, and whether it holds stale + * partitions. + */ + private void recordAssignment(IndexedGroup group, int member, int cohort, int[] subscribed) { + // forEach walks the table of a HashMap directly, measurably cheaper than its iterator on + // large groups. + assignments.get(member).forEach((topicId, partitions) -> { + int topic = group.topicIndex(topicId); + if (topic < 0 || subscribed[topic] != cohort) { + holdsStalePartitions[member] = true; + } else { + recordOwner(group, topic, member, partitions); + } + }); + } + + /** + * Records the member as the owner of its partitions of the topic. + */ + private void recordOwner(IndexedGroup group, int topic, int member, Set<Integer> partitions) { + var topicOwners = owners[topic]; + if (topicOwners == null) { + topicOwners = new int[group.partitionCount(topic)]; + Arrays.fill(topicOwners, -1); + owners[topic] = topicOwners; + } + for (int partition : partitions) { + topicOwners[partition] = member; + } + } + + /** + * @return The current assignment of the member, as given by the spec. + */ + Map<Uuid, Set<Integer>> assignment(int member) { + return assignments.get(member); + } + + /** + * @return {@code true} if the member holds stale partitions. + */ + boolean holdsStalePartitions(int member) { + return holdsStalePartitions[member]; + } + + /** + * @return Per partition of the topic, its owner, or -1 if it has none; or null if no + * partition of the topic has an owner. The array must not be changed. + */ + int[] owners(int topic) { + return owners[topic]; + } + + /** + * Counts the partitions of the topic that each of its owners owns. + * + * @param topic The topic. + * @param counts Receives, at the index of every owner, its number of partitions. Must + * be sized to the number of members, with zero entries; the caller resets + * the entries of the owners to zero after use. Review Comment: nit: I'm still not happy with the wording. `zero` reads like the number of entries rather than describing the contents. How about "zeroed"? And then we can drop the last part about the caller since the caller's obligation is fully pinned down now. ```suggestion * @param counts Receives, at the index of every owner, its number of partitions. Must * be sized to the number of members, with zeroed entries. ``` ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,506 @@ +/* + * 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 the assignment + * is deterministic and does not depend on the iteration order of the collections in the spec. + * + * <p>A cohort is a set of members having the same subscription. Cohorts are numbered in the order + * of their first member. Grouping the members by subscription lets the assignor read, store and + * process every subscription once, however many members share it. + * + * <p>Only topics subscribed by at least one member and which exist are retained. + * + * <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 looked up by a hash of its topics that does not + * depend on their order, see {@link CohortsBySubscription}. The cohorts with the same hash are + * compared exactly, so a collision only costs a comparison, never a wrong cohort. Review Comment: nit, claudish: ```suggestion * compared exactly, so members are never assigned to the wrong cohort. ``` ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,506 @@ +/* + * 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 the assignment + * is deterministic and does not depend on the iteration order of the collections in the spec. + * + * <p>A cohort is a set of members having the same subscription. Cohorts are numbered in the order + * of their first member. Grouping the members by subscription lets the assignor read, store and + * process every subscription once, however many members share it. + * + * <p>Only topics subscribed by at least one member and which exist are retained. + * + * <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 looked up by a hash of its topics that does not + * depend on their order, see {@link CohortsBySubscription}. The cohorts with the same hash are + * compared exactly, so a collision only costs a comparison, never a wrong cohort. + */ +final class IndexedGroup { + /** + * Reads the subscriptions of the members, indexing their topics that exist in the order they + * are seen, and remembering their partition counts. + */ + private static final class SubscriptionReader { + /** + * The group spec, which gives the subscription of every member. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics, which gives their partition counts. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * Per topic seen, its index in the order seen, or -1 if it does not exist, so that the + * describer is asked once per topic. + */ + private final Map<Uuid, Integer> topicIndices = new HashMap<>(); + + /** + * The topics seen that exist, in the order seen. + */ + private final List<Uuid> topicIds = new ArrayList<>(); + + /** + * The partition counts of the topics seen that exist, in the order seen. + */ + private final IntList partitionCounts = new IntList(16); + + SubscriptionReader(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the subscription of the member. + * + * @param memberId The member id. + * @param topics Receives the indices of its subscribed topics that exist. + */ + void readSubscribedTopics(String memberId, IntList topics) { + // Only the iteration is used: the subscription may be a TopicIds, whose size() also + // counts the names that do not resolve. + for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) { + var index = topicIndices.get(topicId); + if (index == null) { + // Seen for the first time: assigned the next index, unless it does not exist. + int partitionCount = subscribedTopicDescriber.numPartitions(topicId); + index = partitionCount < 0 ? -1 : topicIds.size(); + if (index >= 0) { + topicIds.add(topicId); + partitionCounts.add(partitionCount); + } + topicIndices.put(topicId, index); + } + if (index >= 0) { + topics.add(index); + } + } + } + } + + /** + * A hash map from subscriptions to cohorts: {@link #getOrCreateCohort} returns the cohort with + * the exact same subscribed topics, creating it when there is none. + * + * <p>A subscription is summarized by a hash that does not depend on the order of its topics, + * and the cohorts with the same hash are listed together, usually one per hash. + * + * <p>Subscription equality is tested by marking the topics of the subscription being looked + * up in an array, then checking that a cohort has as many topics, all of them marked. Unlike a + * {@code HashMap} keyed by the subscription, this needs no sorted copy of every member's + * subscription. + */ + private static final class CohortsBySubscription { + /** + * Per cohort, its subscribed topics, indexed in the order they were seen. + */ + private final List<int[]> subscriptions = new ArrayList<>(); + + /** + * Per hash of a subscription, the cohorts whose subscription has that hash. + */ + private final Map<Long, IntList> cohortsByHash = new HashMap<>(); + + /** + * Per topic, the lookup that marked it last, or 0 if none did. + */ + private final IntList marks = new IntList(16); + + /** + * The number of lookups so far, which is also the mark of the current one. + */ + private int lookups; + + /** + * @param subscribedTopics The subscribed topics of a member, in any order, without + * duplicates. + * @return The cohort with the exact same subscription, created if there is none. + */ + int getOrCreateCohort(IntList subscribedTopics) { + int mark = ++lookups; + markTopics(subscribedTopics, mark); + long hash = hashOf(subscribedTopics); + var cohorts = cohortsByHash.get(hash); + if (cohorts == null) { + cohorts = new IntList(1); + cohortsByHash.put(hash, cohorts); + } + for (int i = 0; i < cohorts.size(); i++) { + int cohort = cohorts.get(i); + if (hasSameSubscription(cohort, subscribedTopics.size(), mark)) { + return cohort; + } + } + int cohort = subscriptions.size(); + subscriptions.add(subscribedTopics.toArray()); + cohorts.add(cohort); + return cohort; + } + + /** + * @return Per cohort, its subscribed topics. The list must not be changed. + */ + List<int[]> subscriptions() { + return subscriptions; + } + + /** + * Marks the topics with the mark, growing the marks to cover them. + */ + private void markTopics(IntList subscribedTopics, int mark) { + for (int i = 0; i < subscribedTopics.size(); i++) { + int topic = subscribedTopics.get(i); + while (marks.size() <= topic) { + marks.add(0); + } + marks.set(topic, mark); + } + } + + /** + * @return {@code true} if the cohort has the exact same subscription as the one marked with + * the mark, which has that many topics. + */ + private boolean hasSameSubscription(int cohort, int topicCount, int mark) { + int[] topics = subscriptions.get(cohort); + if (topics.length != topicCount) { + return false; + } + // Its topics were marked when it was created, so the marks cover them. + for (int topic : topics) { + if (marks.get(topic) != mark) { + return false; + } + } + return true; + } + + /** + * @return A hash of the topics that does not depend on their order: their number plus the + * sum of a well spread hash of every topic, so that two different subscriptions + * rarely have the same hash. + */ + private static long hashOf(IntList topics) { + long hash = topics.size(); + for (int i = 0; i < topics.size(); i++) { + hash += mix(topics.get(i)); + } + return hash; + } + + /** + * @return A well spread hash of the topic: the SplitMix64 finalizer of its index plus one. + */ + private static long mix(int topic) { + long z = (topic + 1) * 0x9E3779B97F4A7C15L; + z = (z ^ (z >>> 30)) * 0xBF58476D1CE4E5B9L; + z = (z ^ (z >>> 27)) * 0x94D049BB133111EBL; + return z ^ (z >>> 31); + } + } + + /** + * The member ids, in member id order: the id of every member. + */ + private final String[] memberIds; + + /** + * Per member, its cohort. + */ + private final int[] cohortOfMember; + + /** + * The topic ids, in topic id order: the id of every topic. + */ + private final Uuid[] topicIds; + + /** + * Per topic id seen in a subscription, its index, or -1 if the topic does not exist. + */ + private final Map<Uuid, Integer> topicIndices; + + /** + * Per topic, its number of partitions. + */ + private final int[] partitionCounts; + + /** + * Per topic, its number of subscribers. + */ + private final int[] subscriberCounts; + + /** + * Per topic, the cohorts subscribing to it, in cohort order: an array of cohorts, as long as + * the number of such cohorts. + */ + private final int[][] cohortsOfTopic; + + /** + * Per cohort, its members, in member order: an array of members, as long as the cohort. + */ + private final int[][] cohortMembers; + + /** + * Per cohort, its subscribed topics, in topic order: an array of topics, as long as the + * subscription. + */ + private final int[][] cohortTopics; Review Comment: I think it's okay to keep it short here. We don't need to spell out the length or array-ness since it's already implied by the first part of the javadoc. ```suggestion /** * Per topic, the list of cohorts subscribing to it, in cohort order. */ private final int[][] cohortsOfTopic; /** * Per cohort, its members, in member order. */ private final int[][] cohortMembers; /** * Per cohort, its subscribed topics, in topic order. */ private final int[][] cohortTopics; ``` ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,506 @@ +/* + * 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 the assignment + * is deterministic and does not depend on the iteration order of the collections in the spec. + * + * <p>A cohort is a set of members having the same subscription. Cohorts are numbered in the order + * of their first member. Grouping the members by subscription lets the assignor read, store and + * process every subscription once, however many members share it. + * + * <p>Only topics subscribed by at least one member and which exist are retained. + * + * <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 looked up by a hash of its topics that does not + * depend on their order, see {@link CohortsBySubscription}. The cohorts with the same hash are + * compared exactly, so a collision only costs a comparison, never a wrong cohort. + */ +final class IndexedGroup { + /** + * Reads the subscriptions of the members, indexing their topics that exist in the order they + * are seen, and remembering their partition counts. + */ + private static final class SubscriptionReader { + /** + * The group spec, which gives the subscription of every member. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics, which gives their partition counts. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * Per topic seen, its index in the order seen, or -1 if it does not exist, so that the + * describer is asked once per topic. + */ + private final Map<Uuid, Integer> topicIndices = new HashMap<>(); + + /** + * The topics seen that exist, in the order seen. + */ + private final List<Uuid> topicIds = new ArrayList<>(); + + /** + * The partition counts of the topics seen that exist, in the order seen. + */ + private final IntList partitionCounts = new IntList(16); + + SubscriptionReader(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the subscription of the member. + * + * @param memberId The member id. + * @param topics Receives the indices of its subscribed topics that exist. + */ + void readSubscribedTopics(String memberId, IntList topics) { + // Only the iteration is used: the subscription may be a TopicIds, whose size() also + // counts the names that do not resolve. + for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) { + var index = topicIndices.get(topicId); + if (index == null) { + // Seen for the first time: assigned the next index, unless it does not exist. + int partitionCount = subscribedTopicDescriber.numPartitions(topicId); + index = partitionCount < 0 ? -1 : topicIds.size(); + if (index >= 0) { + topicIds.add(topicId); + partitionCounts.add(partitionCount); + } + topicIndices.put(topicId, index); + } + if (index >= 0) { + topics.add(index); + } + } + } + } + + /** + * A hash map from subscriptions to cohorts: {@link #getOrCreateCohort} returns the cohort with + * the exact same subscribed topics, creating it when there is none. + * + * <p>A subscription is summarized by a hash that does not depend on the order of its topics, + * and the cohorts with the same hash are listed together, usually one per hash. + * + * <p>Subscription equality is tested by marking the topics of the subscription being looked + * up in an array, then checking that a cohort has as many topics, all of them marked. Unlike a + * {@code HashMap} keyed by the subscription, this needs no sorted copy of every member's + * subscription. + */ + private static final class CohortsBySubscription { + /** + * Per cohort, its subscribed topics, indexed in the order they were seen. + */ + private final List<int[]> subscriptions = new ArrayList<>(); + + /** + * Per hash of a subscription, the cohorts whose subscription has that hash. + */ + private final Map<Long, IntList> cohortsByHash = new HashMap<>(); + + /** + * Per topic, the lookup that marked it last, or 0 if none did. + */ + private final IntList marks = new IntList(16); + + /** + * The number of lookups so far, which is also the mark of the current one. + */ + private int lookups; + + /** + * @param subscribedTopics The subscribed topics of a member, in any order, without + * duplicates. + * @return The cohort with the exact same subscription, created if there is none. + */ + int getOrCreateCohort(IntList subscribedTopics) { + int mark = ++lookups; + markTopics(subscribedTopics, mark); + long hash = hashOf(subscribedTopics); + var cohorts = cohortsByHash.get(hash); + if (cohorts == null) { + cohorts = new IntList(1); + cohortsByHash.put(hash, cohorts); + } + for (int i = 0; i < cohorts.size(); i++) { + int cohort = cohorts.get(i); + if (hasSameSubscription(cohort, subscribedTopics.size(), mark)) { + return cohort; + } + } + int cohort = subscriptions.size(); + subscriptions.add(subscribedTopics.toArray()); + cohorts.add(cohort); + return cohort; + } + + /** + * @return Per cohort, its subscribed topics. The list must not be changed. + */ + List<int[]> subscriptions() { + return subscriptions; + } + + /** + * Marks the topics with the mark, growing the marks to cover them. Review Comment: nit: Reads a little strangely but I haven't got a better suggestion right now. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,506 @@ +/* + * 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 the assignment + * is deterministic and does not depend on the iteration order of the collections in the spec. + * + * <p>A cohort is a set of members having the same subscription. Cohorts are numbered in the order + * of their first member. Grouping the members by subscription lets the assignor read, store and + * process every subscription once, however many members share it. + * + * <p>Only topics subscribed by at least one member and which exist are retained. + * + * <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 looked up by a hash of its topics that does not + * depend on their order, see {@link CohortsBySubscription}. The cohorts with the same hash are + * compared exactly, so a collision only costs a comparison, never a wrong cohort. + */ +final class IndexedGroup { + /** + * Reads the subscriptions of the members, indexing their topics that exist in the order they + * are seen, and remembering their partition counts. + */ + private static final class SubscriptionReader { + /** + * The group spec, which gives the subscription of every member. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics, which gives their partition counts. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * Per topic seen, its index in the order seen, or -1 if it does not exist, so that the + * describer is asked once per topic. + */ + private final Map<Uuid, Integer> topicIndices = new HashMap<>(); + + /** + * The topics seen that exist, in the order seen. + */ + private final List<Uuid> topicIds = new ArrayList<>(); + + /** + * The partition counts of the topics seen that exist, in the order seen. + */ + private final IntList partitionCounts = new IntList(16); + + SubscriptionReader(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the subscription of the member. + * + * @param memberId The member id. + * @param topics Receives the indices of its subscribed topics that exist. + */ + void readSubscribedTopics(String memberId, IntList topics) { + // Only the iteration is used: the subscription may be a TopicIds, whose size() also + // counts the names that do not resolve. + for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) { + var index = topicIndices.get(topicId); + if (index == null) { + // Seen for the first time: assigned the next index, unless it does not exist. + int partitionCount = subscribedTopicDescriber.numPartitions(topicId); + index = partitionCount < 0 ? -1 : topicIds.size(); + if (index >= 0) { + topicIds.add(topicId); + partitionCounts.add(partitionCount); + } + topicIndices.put(topicId, index); + } + if (index >= 0) { + topics.add(index); + } + } + } + } + + /** + * A hash map from subscriptions to cohorts: {@link #getOrCreateCohort} returns the cohort with + * the exact same subscribed topics, creating it when there is none. Review Comment: ```suggestion * A hash map from subscriptions to cohorts. ``` This part belongs on the javadoc for `getOrCreateCohort`. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,506 @@ +/* + * 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 the assignment + * is deterministic and does not depend on the iteration order of the collections in the spec. + * + * <p>A cohort is a set of members having the same subscription. Cohorts are numbered in the order + * of their first member. Grouping the members by subscription lets the assignor read, store and + * process every subscription once, however many members share it. + * + * <p>Only topics subscribed by at least one member and which exist are retained. + * + * <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 looked up by a hash of its topics that does not + * depend on their order, see {@link CohortsBySubscription}. The cohorts with the same hash are + * compared exactly, so a collision only costs a comparison, never a wrong cohort. + */ +final class IndexedGroup { + /** + * Reads the subscriptions of the members, indexing their topics that exist in the order they + * are seen, and remembering their partition counts. + */ + private static final class SubscriptionReader { + /** + * The group spec, which gives the subscription of every member. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics, which gives their partition counts. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * Per topic seen, its index in the order seen, or -1 if it does not exist, so that the + * describer is asked once per topic. + */ + private final Map<Uuid, Integer> topicIndices = new HashMap<>(); + + /** + * The topics seen that exist, in the order seen. + */ + private final List<Uuid> topicIds = new ArrayList<>(); + + /** + * The partition counts of the topics seen that exist, in the order seen. + */ + private final IntList partitionCounts = new IntList(16); + + SubscriptionReader(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the subscription of the member. + * + * @param memberId The member id. + * @param topics Receives the indices of its subscribed topics that exist. + */ + void readSubscribedTopics(String memberId, IntList topics) { + // Only the iteration is used: the subscription may be a TopicIds, whose size() also + // counts the names that do not resolve. + for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) { + var index = topicIndices.get(topicId); + if (index == null) { + // Seen for the first time: assigned the next index, unless it does not exist. + int partitionCount = subscribedTopicDescriber.numPartitions(topicId); + index = partitionCount < 0 ? -1 : topicIds.size(); + if (index >= 0) { + topicIds.add(topicId); + partitionCounts.add(partitionCount); + } + topicIndices.put(topicId, index); + } + if (index >= 0) { + topics.add(index); + } + } + } + } + + /** + * A hash map from subscriptions to cohorts: {@link #getOrCreateCohort} returns the cohort with + * the exact same subscribed topics, creating it when there is none. + * + * <p>A subscription is summarized by a hash that does not depend on the order of its topics, + * and the cohorts with the same hash are listed together, usually one per hash. + * + * <p>Subscription equality is tested by marking the topics of the subscription being looked + * up in an array, then checking that a cohort has as many topics, all of them marked. Unlike a + * {@code HashMap} keyed by the subscription, this needs no sorted copy of every member's + * subscription. + */ + private static final class CohortsBySubscription { + /** + * Per cohort, its subscribed topics, indexed in the order they were seen. + */ + private final List<int[]> subscriptions = new ArrayList<>(); + + /** + * Per hash of a subscription, the cohorts whose subscription has that hash. + */ + private final Map<Long, IntList> cohortsByHash = new HashMap<>(); + + /** + * Per topic, the lookup that marked it last, or 0 if none did. + */ + private final IntList marks = new IntList(16); + + /** + * The number of lookups so far, which is also the mark of the current one. + */ + private int lookups; + + /** + * @param subscribedTopics The subscribed topics of a member, in any order, without + * duplicates. + * @return The cohort with the exact same subscription, created if there is none. + */ + int getOrCreateCohort(IntList subscribedTopics) { + int mark = ++lookups; + markTopics(subscribedTopics, mark); + long hash = hashOf(subscribedTopics); + var cohorts = cohortsByHash.get(hash); + if (cohorts == null) { + cohorts = new IntList(1); + cohortsByHash.put(hash, cohorts); + } + for (int i = 0; i < cohorts.size(); i++) { + int cohort = cohorts.get(i); + if (hasSameSubscription(cohort, subscribedTopics.size(), mark)) { + return cohort; + } + } + int cohort = subscriptions.size(); + subscriptions.add(subscribedTopics.toArray()); + cohorts.add(cohort); + return cohort; + } + + /** + * @return Per cohort, its subscribed topics. The list must not be changed. + */ + List<int[]> subscriptions() { + return subscriptions; + } + + /** + * Marks the topics with the mark, growing the marks to cover them. + */ + private void markTopics(IntList subscribedTopics, int mark) { + for (int i = 0; i < subscribedTopics.size(); i++) { + int topic = subscribedTopics.get(i); + while (marks.size() <= topic) { + marks.add(0); + } + marks.set(topic, mark); + } + } + + /** + * @return {@code true} if the cohort has the exact same subscription as the one marked with + * the mark, which has that many topics. Review Comment: ```suggestion * the mark. ``` `that many` doesn't refer to anything in the sentence? ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,506 @@ +/* + * 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 the assignment + * is deterministic and does not depend on the iteration order of the collections in the spec. + * + * <p>A cohort is a set of members having the same subscription. Cohorts are numbered in the order + * of their first member. Grouping the members by subscription lets the assignor read, store and + * process every subscription once, however many members share it. + * + * <p>Only topics subscribed by at least one member and which exist are retained. + * + * <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 looked up by a hash of its topics that does not + * depend on their order, see {@link CohortsBySubscription}. The cohorts with the same hash are + * compared exactly, so a collision only costs a comparison, never a wrong cohort. + */ +final class IndexedGroup { + /** + * Reads the subscriptions of the members, indexing their topics that exist in the order they + * are seen, and remembering their partition counts. + */ + private static final class SubscriptionReader { + /** + * The group spec, which gives the subscription of every member. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics, which gives their partition counts. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * Per topic seen, its index in the order seen, or -1 if it does not exist, so that the + * describer is asked once per topic. + */ + private final Map<Uuid, Integer> topicIndices = new HashMap<>(); + + /** + * The topics seen that exist, in the order seen. + */ + private final List<Uuid> topicIds = new ArrayList<>(); + + /** + * The partition counts of the topics seen that exist, in the order seen. + */ + private final IntList partitionCounts = new IntList(16); + + SubscriptionReader(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the subscription of the member. + * + * @param memberId The member id. + * @param topics Receives the indices of its subscribed topics that exist. + */ + void readSubscribedTopics(String memberId, IntList topics) { + // Only the iteration is used: the subscription may be a TopicIds, whose size() also + // counts the names that do not resolve. + for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) { + var index = topicIndices.get(topicId); + if (index == null) { + // Seen for the first time: assigned the next index, unless it does not exist. + int partitionCount = subscribedTopicDescriber.numPartitions(topicId); + index = partitionCount < 0 ? -1 : topicIds.size(); + if (index >= 0) { + topicIds.add(topicId); + partitionCounts.add(partitionCount); + } + topicIndices.put(topicId, index); + } + if (index >= 0) { + topics.add(index); + } + } + } + } + + /** + * A hash map from subscriptions to cohorts: {@link #getOrCreateCohort} returns the cohort with + * the exact same subscribed topics, creating it when there is none. + * + * <p>A subscription is summarized by a hash that does not depend on the order of its topics, + * and the cohorts with the same hash are listed together, usually one per hash. + * + * <p>Subscription equality is tested by marking the topics of the subscription being looked + * up in an array, then checking that a cohort has as many topics, all of them marked. Unlike a + * {@code HashMap} keyed by the subscription, this needs no sorted copy of every member's + * subscription. + */ + private static final class CohortsBySubscription { + /** + * Per cohort, its subscribed topics, indexed in the order they were seen. + */ + private final List<int[]> subscriptions = new ArrayList<>(); + + /** + * Per hash of a subscription, the cohorts whose subscription has that hash. + */ + private final Map<Long, IntList> cohortsByHash = new HashMap<>(); + + /** + * Per topic, the lookup that marked it last, or 0 if none did. + */ + private final IntList marks = new IntList(16); + + /** + * The number of lookups so far, which is also the mark of the current one. + */ + private int lookups; + + /** + * @param subscribedTopics The subscribed topics of a member, in any order, without + * duplicates. + * @return The cohort with the exact same subscription, created if there is none. + */ + int getOrCreateCohort(IntList subscribedTopics) { + int mark = ++lookups; + markTopics(subscribedTopics, mark); + long hash = hashOf(subscribedTopics); + var cohorts = cohortsByHash.get(hash); + if (cohorts == null) { + cohorts = new IntList(1); + cohortsByHash.put(hash, cohorts); + } + for (int i = 0; i < cohorts.size(); i++) { + int cohort = cohorts.get(i); + if (hasSameSubscription(cohort, subscribedTopics.size(), mark)) { + return cohort; + } + } + int cohort = subscriptions.size(); + subscriptions.add(subscribedTopics.toArray()); + cohorts.add(cohort); + return cohort; + } + + /** + * @return Per cohort, its subscribed topics. The list must not be changed. + */ + List<int[]> subscriptions() { + return subscriptions; + } + + /** + * Marks the topics with the mark, growing the marks to cover them. + */ + private void markTopics(IntList subscribedTopics, int mark) { + for (int i = 0; i < subscribedTopics.size(); i++) { + int topic = subscribedTopics.get(i); + while (marks.size() <= topic) { + marks.add(0); + } + marks.set(topic, mark); + } + } + + /** + * @return {@code true} if the cohort has the exact same subscription as the one marked with + * the mark, which has that many topics. + */ + private boolean hasSameSubscription(int cohort, int topicCount, int mark) { + int[] topics = subscriptions.get(cohort); + if (topics.length != topicCount) { + return false; + } + // Its topics were marked when it was created, so the marks cover them. + for (int topic : topics) { + if (marks.get(topic) != mark) { + return false; + } + } + return true; + } + + /** + * @return A hash of the topics that does not depend on their order: their number plus the + * sum of a well spread hash of every topic, so that two different subscriptions + * rarely have the same hash. + */ + private static long hashOf(IntList topics) { + long hash = topics.size(); + for (int i = 0; i < topics.size(); i++) { + hash += mix(topics.get(i)); + } + return hash; + } + + /** + * @return A well spread hash of the topic: the SplitMix64 finalizer of its index plus one. + */ + private static long mix(int topic) { Review Comment: Since this is also a hash, shall we call it `hashOf`? Or if we don't want to re-use the name, we can have `topicHash` and `subscriptionHash`? ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,506 @@ +/* + * 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 the assignment + * is deterministic and does not depend on the iteration order of the collections in the spec. + * + * <p>A cohort is a set of members having the same subscription. Cohorts are numbered in the order + * of their first member. Grouping the members by subscription lets the assignor read, store and + * process every subscription once, however many members share it. + * + * <p>Only topics subscribed by at least one member and which exist are retained. + * + * <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 looked up by a hash of its topics that does not + * depend on their order, see {@link CohortsBySubscription}. The cohorts with the same hash are + * compared exactly, so a collision only costs a comparison, never a wrong cohort. + */ +final class IndexedGroup { + /** + * Reads the subscriptions of the members, indexing their topics that exist in the order they + * are seen, and remembering their partition counts. + */ + private static final class SubscriptionReader { + /** + * The group spec, which gives the subscription of every member. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics, which gives their partition counts. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * Per topic seen, its index in the order seen, or -1 if it does not exist, so that the + * describer is asked once per topic. + */ + private final Map<Uuid, Integer> topicIndices = new HashMap<>(); + + /** + * The topics seen that exist, in the order seen. + */ + private final List<Uuid> topicIds = new ArrayList<>(); + + /** + * The partition counts of the topics seen that exist, in the order seen. + */ + private final IntList partitionCounts = new IntList(16); + + SubscriptionReader(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the subscription of the member. + * + * @param memberId The member id. + * @param topics Receives the indices of its subscribed topics that exist. + */ + void readSubscribedTopics(String memberId, IntList topics) { + // Only the iteration is used: the subscription may be a TopicIds, whose size() also + // counts the names that do not resolve. + for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) { + var index = topicIndices.get(topicId); + if (index == null) { + // Seen for the first time: assigned the next index, unless it does not exist. + int partitionCount = subscribedTopicDescriber.numPartitions(topicId); + index = partitionCount < 0 ? -1 : topicIds.size(); + if (index >= 0) { + topicIds.add(topicId); + partitionCounts.add(partitionCount); + } + topicIndices.put(topicId, index); + } + if (index >= 0) { + topics.add(index); + } + } + } + } + + /** + * A hash map from subscriptions to cohorts: {@link #getOrCreateCohort} returns the cohort with + * the exact same subscribed topics, creating it when there is none. + * + * <p>A subscription is summarized by a hash that does not depend on the order of its topics, + * and the cohorts with the same hash are listed together, usually one per hash. + * + * <p>Subscription equality is tested by marking the topics of the subscription being looked + * up in an array, then checking that a cohort has as many topics, all of them marked. Unlike a + * {@code HashMap} keyed by the subscription, this needs no sorted copy of every member's + * subscription. + */ + private static final class CohortsBySubscription { + /** + * Per cohort, its subscribed topics, indexed in the order they were seen. + */ + private final List<int[]> subscriptions = new ArrayList<>(); + + /** + * Per hash of a subscription, the cohorts whose subscription has that hash. + */ + private final Map<Long, IntList> cohortsByHash = new HashMap<>(); + + /** + * Per topic, the lookup that marked it last, or 0 if none did. + */ + private final IntList marks = new IntList(16); + + /** + * The number of lookups so far, which is also the mark of the current one. + */ + private int lookups; + + /** + * @param subscribedTopics The subscribed topics of a member, in any order, without + * duplicates. + * @return The cohort with the exact same subscription, created if there is none. + */ + int getOrCreateCohort(IntList subscribedTopics) { + int mark = ++lookups; + markTopics(subscribedTopics, mark); + long hash = hashOf(subscribedTopics); + var cohorts = cohortsByHash.get(hash); + if (cohorts == null) { + cohorts = new IntList(1); + cohortsByHash.put(hash, cohorts); + } + for (int i = 0; i < cohorts.size(); i++) { + int cohort = cohorts.get(i); + if (hasSameSubscription(cohort, subscribedTopics.size(), mark)) { + return cohort; + } + } + int cohort = subscriptions.size(); + subscriptions.add(subscribedTopics.toArray()); + cohorts.add(cohort); + return cohort; + } + + /** + * @return Per cohort, its subscribed topics. The list must not be changed. + */ + List<int[]> subscriptions() { + return subscriptions; + } + + /** + * Marks the topics with the mark, growing the marks to cover them. + */ + private void markTopics(IntList subscribedTopics, int mark) { + for (int i = 0; i < subscribedTopics.size(); i++) { + int topic = subscribedTopics.get(i); + while (marks.size() <= topic) { + marks.add(0); + } + marks.set(topic, mark); + } + } + + /** + * @return {@code true} if the cohort has the exact same subscription as the one marked with + * the mark, which has that many topics. + */ + private boolean hasSameSubscription(int cohort, int topicCount, int mark) { + int[] topics = subscriptions.get(cohort); + if (topics.length != topicCount) { + return false; + } + // Its topics were marked when it was created, so the marks cover them. + for (int topic : topics) { + if (marks.get(topic) != mark) { + return false; + } + } + return true; + } + + /** + * @return A hash of the topics that does not depend on their order: their number plus the + * sum of a well spread hash of every topic, so that two different subscriptions + * rarely have the same hash. + */ + private static long hashOf(IntList topics) { + long hash = topics.size(); + for (int i = 0; i < topics.size(); i++) { + hash += mix(topics.get(i)); + } + return hash; + } + + /** + * @return A well spread hash of the topic: the SplitMix64 finalizer of its index plus one. + */ + private static long mix(int topic) { + long z = (topic + 1) * 0x9E3779B97F4A7C15L; + z = (z ^ (z >>> 30)) * 0xBF58476D1CE4E5B9L; + z = (z ^ (z >>> 27)) * 0x94D049BB133111EBL; + return z ^ (z >>> 31); + } + } + + /** + * The member ids, in member id order: the id of every member. Review Comment: We can drop the last part. It's the same words in a different order ```suggestion * The member ids, in member id order. ``` or to avoid repetition: ```suggestion * The member ids, in ascending order. ``` ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,506 @@ +/* + * 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 the assignment + * is deterministic and does not depend on the iteration order of the collections in the spec. + * + * <p>A cohort is a set of members having the same subscription. Cohorts are numbered in the order + * of their first member. Grouping the members by subscription lets the assignor read, store and + * process every subscription once, however many members share it. + * + * <p>Only topics subscribed by at least one member and which exist are retained. + * + * <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 looked up by a hash of its topics that does not + * depend on their order, see {@link CohortsBySubscription}. The cohorts with the same hash are + * compared exactly, so a collision only costs a comparison, never a wrong cohort. + */ +final class IndexedGroup { + /** + * Reads the subscriptions of the members, indexing their topics that exist in the order they + * are seen, and remembering their partition counts. + */ + private static final class SubscriptionReader { + /** + * The group spec, which gives the subscription of every member. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics, which gives their partition counts. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * Per topic seen, its index in the order seen, or -1 if it does not exist, so that the + * describer is asked once per topic. + */ + private final Map<Uuid, Integer> topicIndices = new HashMap<>(); + + /** + * The topics seen that exist, in the order seen. + */ + private final List<Uuid> topicIds = new ArrayList<>(); + + /** + * The partition counts of the topics seen that exist, in the order seen. + */ + private final IntList partitionCounts = new IntList(16); + + SubscriptionReader(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the subscription of the member. + * + * @param memberId The member id. + * @param topics Receives the indices of its subscribed topics that exist. + */ + void readSubscribedTopics(String memberId, IntList topics) { + // Only the iteration is used: the subscription may be a TopicIds, whose size() also + // counts the names that do not resolve. + for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) { + var index = topicIndices.get(topicId); + if (index == null) { + // Seen for the first time: assigned the next index, unless it does not exist. + int partitionCount = subscribedTopicDescriber.numPartitions(topicId); + index = partitionCount < 0 ? -1 : topicIds.size(); + if (index >= 0) { + topicIds.add(topicId); + partitionCounts.add(partitionCount); + } + topicIndices.put(topicId, index); + } + if (index >= 0) { + topics.add(index); + } + } + } + } + + /** + * A hash map from subscriptions to cohorts: {@link #getOrCreateCohort} returns the cohort with + * the exact same subscribed topics, creating it when there is none. + * + * <p>A subscription is summarized by a hash that does not depend on the order of its topics, + * and the cohorts with the same hash are listed together, usually one per hash. + * + * <p>Subscription equality is tested by marking the topics of the subscription being looked + * up in an array, then checking that a cohort has as many topics, all of them marked. Unlike a + * {@code HashMap} keyed by the subscription, this needs no sorted copy of every member's + * subscription. + */ + private static final class CohortsBySubscription { + /** + * Per cohort, its subscribed topics, indexed in the order they were seen. + */ + private final List<int[]> subscriptions = new ArrayList<>(); + + /** + * Per hash of a subscription, the cohorts whose subscription has that hash. + */ + private final Map<Long, IntList> cohortsByHash = new HashMap<>(); + + /** + * Per topic, the lookup that marked it last, or 0 if none did. + */ + private final IntList marks = new IntList(16); + + /** + * The number of lookups so far, which is also the mark of the current one. + */ + private int lookups; + + /** + * @param subscribedTopics The subscribed topics of a member, in any order, without + * duplicates. + * @return The cohort with the exact same subscription, created if there is none. + */ + int getOrCreateCohort(IntList subscribedTopics) { + int mark = ++lookups; + markTopics(subscribedTopics, mark); + long hash = hashOf(subscribedTopics); + var cohorts = cohortsByHash.get(hash); + if (cohorts == null) { + cohorts = new IntList(1); + cohortsByHash.put(hash, cohorts); + } + for (int i = 0; i < cohorts.size(); i++) { + int cohort = cohorts.get(i); + if (hasSameSubscription(cohort, subscribedTopics.size(), mark)) { + return cohort; + } + } + int cohort = subscriptions.size(); + subscriptions.add(subscribedTopics.toArray()); + cohorts.add(cohort); + return cohort; + } + + /** + * @return Per cohort, its subscribed topics. The list must not be changed. + */ + List<int[]> subscriptions() { + return subscriptions; + } + + /** + * Marks the topics with the mark, growing the marks to cover them. + */ + private void markTopics(IntList subscribedTopics, int mark) { + for (int i = 0; i < subscribedTopics.size(); i++) { + int topic = subscribedTopics.get(i); + while (marks.size() <= topic) { + marks.add(0); + } + marks.set(topic, mark); + } + } + + /** + * @return {@code true} if the cohort has the exact same subscription as the one marked with + * the mark, which has that many topics. + */ + private boolean hasSameSubscription(int cohort, int topicCount, int mark) { + int[] topics = subscriptions.get(cohort); + if (topics.length != topicCount) { + return false; + } + // Its topics were marked when it was created, so the marks cover them. Review Comment: I can't tell what `its` refers to here. ```suggestion ``` ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,506 @@ +/* + * 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 the assignment + * is deterministic and does not depend on the iteration order of the collections in the spec. + * + * <p>A cohort is a set of members having the same subscription. Cohorts are numbered in the order + * of their first member. Grouping the members by subscription lets the assignor read, store and + * process every subscription once, however many members share it. + * + * <p>Only topics subscribed by at least one member and which exist are retained. + * + * <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 looked up by a hash of its topics that does not + * depend on their order, see {@link CohortsBySubscription}. The cohorts with the same hash are + * compared exactly, so a collision only costs a comparison, never a wrong cohort. + */ +final class IndexedGroup { + /** + * Reads the subscriptions of the members, indexing their topics that exist in the order they + * are seen, and remembering their partition counts. + */ + private static final class SubscriptionReader { + /** + * The group spec, which gives the subscription of every member. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics, which gives their partition counts. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * Per topic seen, its index in the order seen, or -1 if it does not exist, so that the + * describer is asked once per topic. + */ + private final Map<Uuid, Integer> topicIndices = new HashMap<>(); + + /** + * The topics seen that exist, in the order seen. + */ + private final List<Uuid> topicIds = new ArrayList<>(); + + /** + * The partition counts of the topics seen that exist, in the order seen. + */ + private final IntList partitionCounts = new IntList(16); + + SubscriptionReader(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the subscription of the member. + * + * @param memberId The member id. + * @param topics Receives the indices of its subscribed topics that exist. + */ + void readSubscribedTopics(String memberId, IntList topics) { + // Only the iteration is used: the subscription may be a TopicIds, whose size() also + // counts the names that do not resolve. + for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) { + var index = topicIndices.get(topicId); + if (index == null) { + // Seen for the first time: assigned the next index, unless it does not exist. + int partitionCount = subscribedTopicDescriber.numPartitions(topicId); + index = partitionCount < 0 ? -1 : topicIds.size(); + if (index >= 0) { + topicIds.add(topicId); + partitionCounts.add(partitionCount); + } + topicIndices.put(topicId, index); + } + if (index >= 0) { + topics.add(index); + } + } + } + } + + /** + * A hash map from subscriptions to cohorts: {@link #getOrCreateCohort} returns the cohort with + * the exact same subscribed topics, creating it when there is none. + * + * <p>A subscription is summarized by a hash that does not depend on the order of its topics, + * and the cohorts with the same hash are listed together, usually one per hash. + * + * <p>Subscription equality is tested by marking the topics of the subscription being looked + * up in an array, then checking that a cohort has as many topics, all of them marked. Unlike a + * {@code HashMap} keyed by the subscription, this needs no sorted copy of every member's + * subscription. + */ + private static final class CohortsBySubscription { + /** + * Per cohort, its subscribed topics, indexed in the order they were seen. + */ + private final List<int[]> subscriptions = new ArrayList<>(); + + /** + * Per hash of a subscription, the cohorts whose subscription has that hash. + */ + private final Map<Long, IntList> cohortsByHash = new HashMap<>(); + + /** + * Per topic, the lookup that marked it last, or 0 if none did. + */ + private final IntList marks = new IntList(16); + + /** + * The number of lookups so far, which is also the mark of the current one. + */ + private int lookups; + + /** + * @param subscribedTopics The subscribed topics of a member, in any order, without + * duplicates. + * @return The cohort with the exact same subscription, created if there is none. + */ + int getOrCreateCohort(IntList subscribedTopics) { + int mark = ++lookups; + markTopics(subscribedTopics, mark); + long hash = hashOf(subscribedTopics); + var cohorts = cohortsByHash.get(hash); + if (cohorts == null) { + cohorts = new IntList(1); + cohortsByHash.put(hash, cohorts); + } + for (int i = 0; i < cohorts.size(); i++) { + int cohort = cohorts.get(i); + if (hasSameSubscription(cohort, subscribedTopics.size(), mark)) { + return cohort; + } + } + int cohort = subscriptions.size(); + subscriptions.add(subscribedTopics.toArray()); + cohorts.add(cohort); + return cohort; + } + + /** + * @return Per cohort, its subscribed topics. The list must not be changed. + */ + List<int[]> subscriptions() { + return subscriptions; + } + + /** + * Marks the topics with the mark, growing the marks to cover them. + */ + private void markTopics(IntList subscribedTopics, int mark) { + for (int i = 0; i < subscribedTopics.size(); i++) { + int topic = subscribedTopics.get(i); + while (marks.size() <= topic) { + marks.add(0); + } + marks.set(topic, mark); + } + } + + /** + * @return {@code true} if the cohort has the exact same subscription as the one marked with + * the mark, which has that many topics. + */ + private boolean hasSameSubscription(int cohort, int topicCount, int mark) { + int[] topics = subscriptions.get(cohort); + if (topics.length != topicCount) { + return false; + } + // Its topics were marked when it was created, so the marks cover them. + for (int topic : topics) { + if (marks.get(topic) != mark) { + return false; + } + } + return true; + } + + /** + * @return A hash of the topics that does not depend on their order: their number plus the + * sum of a well spread hash of every topic, so that two different subscriptions + * rarely have the same hash. + */ + private static long hashOf(IntList topics) { + long hash = topics.size(); + for (int i = 0; i < topics.size(); i++) { + hash += mix(topics.get(i)); + } + return hash; + } + + /** + * @return A well spread hash of the topic: the SplitMix64 finalizer of its index plus one. + */ + private static long mix(int topic) { + long z = (topic + 1) * 0x9E3779B97F4A7C15L; + z = (z ^ (z >>> 30)) * 0xBF58476D1CE4E5B9L; + z = (z ^ (z >>> 27)) * 0x94D049BB133111EBL; + return z ^ (z >>> 31); + } + } + + /** + * The member ids, in member id order: the id of every member. + */ + private final String[] memberIds; + + /** + * Per member, its cohort. + */ + private final int[] cohortOfMember; + + /** + * The topic ids, in topic id order: the id of every topic. Review Comment: Same here. ```suggestion * The topic ids, in topic id order. ``` or to avoid repetition: ```suggestion * The topic ids, in ascending order. ``` ########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/TopicsFixture.java: ########## @@ -0,0 +1,61 @@ +/* + * 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.common.runtime.MetadataImageBuilder; +import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber; +import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl; + +/** + * The topics of a test, described from a metadata image as the coordinator describes them. The + * topics are named after their ids. + * + * <p>The describer is built once: the fixture cannot change afterwards. + */ +public final class TopicsFixture { Review Comment: Makes sense ########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroupTest.java: ########## @@ -0,0 +1,272 @@ +/* + * 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.SubscribedTopicDescriber; +import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import java.util.stream.IntStream; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * Tests of {@link IndexedGroup}: the numbering of the members, the numbering of the topics, and the + * cohorts. + * + * <p>The members are added to the specs, and the subscriptions list their topics, out of id order, + * so that the indices cannot come from the iteration order of the spec. + */ +public class IndexedGroupTest { + private static final Uuid T1 = new Uuid(1L, 1L); + private static final Uuid T2 = new Uuid(2L, 1L); + private static final Uuid T3 = new Uuid(3L, 1L); + private static final Uuid MISSING = new Uuid(9L, 1L); + + private static final SubscribedTopicDescriber DESCRIBER = new TopicsFixture() + .withTopic(T1, 2) + .withTopic(T2, 5) + .withTopic(T3, 7) + .build(); + + // Numbering of the members. + + /** + * Member ids are compared as strings, so m10 comes between m1 and m2. + */ + @Test + public void testMembersAreNumberedInMemberIdOrder() { + var spec = new GroupSpecFixture() + .withMember("m2", Set.of(T1)) + .withMember("m10", Set.of(T1)) + .withMember("m1", Set.of(T1)) + .build(); + + var group = new IndexedGroup(spec, DESCRIBER); + + assertEquals(List.of("m1", "m10", "m2"), memberIds(group)); + } + + // Numbering of the topics. + + /** + * A, read first, sees T3 before T1, and B then sees T2; the topics are still numbered in topic + * id order, with their partition counts, and so are the topics of the cohorts. + */ + @Test + public void testTopicsAreNumberedInTopicIdOrder() { + var spec = new GroupSpecFixture() + .withMember("B", Set.of(T2)) + .withMember("A", orderedSet(T3, T1)) + .build(); + + var group = new IndexedGroup(spec, DESCRIBER); + + assertEquals(List.of(T1, T2, T3), topicIds(group)); + assertEquals(List.of(0, 1, 2), topicIndices(group, T1, T2, T3)); + assertEquals(List.of(2, 5, 7), partitionCounts(group)); + assertEquals(List.of(List.of(0, 2), List.of(1)), topicsOfCohorts(group)); + } + + /** + * A subscribes to T1: only T1 has an index, and any other topic gets -1. + */ + @Test + public void testTopicIndex() { + var spec = new GroupSpecFixture() + .withMember("A", Set.of(T1)) + .build(); + + var group = new IndexedGroup(spec, DESCRIBER); + + assertEquals(List.of(T1), topicIds(group)); + assertEquals(List.of(0, -1, -1), topicIndices(group, T1, T2, MISSING)); + assertEquals(List.of(2), partitionCounts(group)); + } + + /** + * The describer gives -1 partitions for a topic that does not exist, which is then dropped from + * the subscriptions: A, subscribing to T1 and the missing topic, and B, subscribing to T1, + * share cohort 0, and C, subscribing only to the missing topic, has a cohort without topics. + */ + @Test + public void testTopicsThatDoNotExistAreIgnored() { + var spec = new GroupSpecFixture() + .withMember("C", Set.of(MISSING)) + .withMember("B", Set.of(T1)) + .withMember("A", orderedSet(MISSING, T1)) + .build(); + + var group = new IndexedGroup(spec, DESCRIBER); + + assertEquals(List.of(T1), topicIds(group)); + assertEquals(List.of(0, -1), topicIndices(group, T1, MISSING)); + assertEquals(List.of(2), partitionCounts(group)); + assertEquals(List.of(List.of(0, 1), List.of(2)), membersOfCohorts(group)); + assertEquals(List.of(List.of(0), List.of()), topicsOfCohorts(group)); Review Comment: Shall we assert on `membersOfCohorts` and `topicsOfCohorts` in the other two topic tests for consistency? ########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroupTest.java: ########## @@ -0,0 +1,272 @@ +/* + * 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.SubscribedTopicDescriber; +import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import java.util.stream.IntStream; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * Tests of {@link IndexedGroup}: the numbering of the members, the numbering of the topics, and the + * cohorts. + * + * <p>The members are added to the specs, and the subscriptions list their topics, out of id order, + * so that the indices cannot come from the iteration order of the spec. + */ +public class IndexedGroupTest { + private static final Uuid T1 = new Uuid(1L, 1L); + private static final Uuid T2 = new Uuid(2L, 1L); + private static final Uuid T3 = new Uuid(3L, 1L); + private static final Uuid MISSING = new Uuid(9L, 1L); + + private static final SubscribedTopicDescriber DESCRIBER = new TopicsFixture() + .withTopic(T1, 2) + .withTopic(T2, 5) + .withTopic(T3, 7) + .build(); + + // Numbering of the members. + + /** + * Member ids are compared as strings, so m10 comes between m1 and m2. + */ + @Test + public void testMembersAreNumberedInMemberIdOrder() { + var spec = new GroupSpecFixture() + .withMember("m2", Set.of(T1)) + .withMember("m10", Set.of(T1)) + .withMember("m1", Set.of(T1)) + .build(); + + var group = new IndexedGroup(spec, DESCRIBER); + + assertEquals(List.of("m1", "m10", "m2"), memberIds(group)); + } + + // Numbering of the topics. + + /** + * A, read first, sees T3 before T1, and B then sees T2; the topics are still numbered in topic + * id order, with their partition counts, and so are the topics of the cohorts. + */ + @Test + public void testTopicsAreNumberedInTopicIdOrder() { + var spec = new GroupSpecFixture() + .withMember("B", Set.of(T2)) + .withMember("A", orderedSet(T3, T1)) + .build(); + + var group = new IndexedGroup(spec, DESCRIBER); + + assertEquals(List.of(T1, T2, T3), topicIds(group)); + assertEquals(List.of(0, 1, 2), topicIndices(group, T1, T2, T3)); + assertEquals(List.of(2, 5, 7), partitionCounts(group)); + assertEquals(List.of(List.of(0, 2), List.of(1)), topicsOfCohorts(group)); + } + + /** + * A subscribes to T1: only T1 has an index, and any other topic gets -1. + */ + @Test + public void testTopicIndex() { Review Comment: Looking at the topic tests again, there's some overlap in the coverage. `testTopicsAreNumberedInTopicIdOrder` already tests topic indexing. Shall we turn this method into `testTopicsThatArentSubscribedAreIgnored`? And drop the index lookup for `MISSING`, since that's the job of `testTopicsThatDoNotExistAreIgnored` to test. Then again, I wonder if this is worth testing at all, since you'd have to really go out of your way to include unsubscribed topics in the numbering. ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,506 @@ +/* + * 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 the assignment + * is deterministic and does not depend on the iteration order of the collections in the spec. + * + * <p>A cohort is a set of members having the same subscription. Cohorts are numbered in the order + * of their first member. Grouping the members by subscription lets the assignor read, store and + * process every subscription once, however many members share it. + * + * <p>Only topics subscribed by at least one member and which exist are retained. + * + * <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 looked up by a hash of its topics that does not + * depend on their order, see {@link CohortsBySubscription}. The cohorts with the same hash are + * compared exactly, so a collision only costs a comparison, never a wrong cohort. + */ +final class IndexedGroup { + /** + * Reads the subscriptions of the members, indexing their topics that exist in the order they + * are seen, and remembering their partition counts. + */ + private static final class SubscriptionReader { + /** + * The group spec, which gives the subscription of every member. + */ + private final GroupSpec groupSpec; + + /** + * The describer of the subscribed topics, which gives their partition counts. + */ + private final SubscribedTopicDescriber subscribedTopicDescriber; + + /** + * Per topic seen, its index in the order seen, or -1 if it does not exist, so that the + * describer is asked once per topic. + */ + private final Map<Uuid, Integer> topicIndices = new HashMap<>(); + + /** + * The topics seen that exist, in the order seen. + */ + private final List<Uuid> topicIds = new ArrayList<>(); + + /** + * The partition counts of the topics seen that exist, in the order seen. + */ + private final IntList partitionCounts = new IntList(16); + + SubscriptionReader(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + this.groupSpec = groupSpec; + this.subscribedTopicDescriber = subscribedTopicDescriber; + } + + /** + * Reads the subscription of the member. + * + * @param memberId The member id. + * @param topics Receives the indices of its subscribed topics that exist. + */ + void readSubscribedTopics(String memberId, IntList topics) { + // Only the iteration is used: the subscription may be a TopicIds, whose size() also + // counts the names that do not resolve. + for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) { + var index = topicIndices.get(topicId); + if (index == null) { + // Seen for the first time: assigned the next index, unless it does not exist. + int partitionCount = subscribedTopicDescriber.numPartitions(topicId); + index = partitionCount < 0 ? -1 : topicIds.size(); + if (index >= 0) { + topicIds.add(topicId); + partitionCounts.add(partitionCount); + } + topicIndices.put(topicId, index); + } + if (index >= 0) { + topics.add(index); + } + } + } + } + + /** + * A hash map from subscriptions to cohorts: {@link #getOrCreateCohort} returns the cohort with + * the exact same subscribed topics, creating it when there is none. + * + * <p>A subscription is summarized by a hash that does not depend on the order of its topics, + * and the cohorts with the same hash are listed together, usually one per hash. + * + * <p>Subscription equality is tested by marking the topics of the subscription being looked + * up in an array, then checking that a cohort has as many topics, all of them marked. Unlike a + * {@code HashMap} keyed by the subscription, this needs no sorted copy of every member's + * subscription. + */ + private static final class CohortsBySubscription { + /** + * Per cohort, its subscribed topics, indexed in the order they were seen. + */ + private final List<int[]> subscriptions = new ArrayList<>(); + + /** + * Per hash of a subscription, the cohorts whose subscription has that hash. + */ + private final Map<Long, IntList> cohortsByHash = new HashMap<>(); + + /** + * Per topic, the lookup that marked it last, or 0 if none did. + */ + private final IntList marks = new IntList(16); + + /** + * The number of lookups so far, which is also the mark of the current one. + */ + private int lookups; + + /** + * @param subscribedTopics The subscribed topics of a member, in any order, without + * duplicates. + * @return The cohort with the exact same subscription, created if there is none. + */ + int getOrCreateCohort(IntList subscribedTopics) { + int mark = ++lookups; + markTopics(subscribedTopics, mark); + long hash = hashOf(subscribedTopics); + var cohorts = cohortsByHash.get(hash); + if (cohorts == null) { + cohorts = new IntList(1); + cohortsByHash.put(hash, cohorts); + } + for (int i = 0; i < cohorts.size(); i++) { + int cohort = cohorts.get(i); + if (hasSameSubscription(cohort, subscribedTopics.size(), mark)) { + return cohort; + } + } + int cohort = subscriptions.size(); + subscriptions.add(subscribedTopics.toArray()); + cohorts.add(cohort); + return cohort; + } + + /** + * @return Per cohort, its subscribed topics. The list must not be changed. + */ + List<int[]> subscriptions() { + return subscriptions; + } + + /** + * Marks the topics with the mark, growing the marks to cover them. + */ + private void markTopics(IntList subscribedTopics, int mark) { + for (int i = 0; i < subscribedTopics.size(); i++) { + int topic = subscribedTopics.get(i); + while (marks.size() <= topic) { + marks.add(0); + } + marks.set(topic, mark); + } + } + + /** + * @return {@code true} if the cohort has the exact same subscription as the one marked with + * the mark, which has that many topics. + */ + private boolean hasSameSubscription(int cohort, int topicCount, int mark) { + int[] topics = subscriptions.get(cohort); + if (topics.length != topicCount) { + return false; + } + // Its topics were marked when it was created, so the marks cover them. + for (int topic : topics) { + if (marks.get(topic) != mark) { + return false; + } + } + return true; + } + + /** + * @return A hash of the topics that does not depend on their order: their number plus the + * sum of a well spread hash of every topic, so that two different subscriptions + * rarely have the same hash. + */ + private static long hashOf(IntList topics) { + long hash = topics.size(); + for (int i = 0; i < topics.size(); i++) { + hash += mix(topics.get(i)); + } + return hash; + } + + /** + * @return A well spread hash of the topic: the SplitMix64 finalizer of its index plus one. + */ + private static long mix(int topic) { + long z = (topic + 1) * 0x9E3779B97F4A7C15L; + z = (z ^ (z >>> 30)) * 0xBF58476D1CE4E5B9L; + z = (z ^ (z >>> 27)) * 0x94D049BB133111EBL; + return z ^ (z >>> 31); + } + } + + /** + * The member ids, in member id order: the id of every member. + */ + private final String[] memberIds; + + /** + * Per member, its cohort. + */ + private final int[] cohortOfMember; + + /** + * The topic ids, in topic id order: the id of every topic. + */ + private final Uuid[] topicIds; + + /** + * Per topic id seen in a subscription, its index, or -1 if the topic does not exist. + */ + private final Map<Uuid, Integer> topicIndices; + + /** + * Per topic, its number of partitions. + */ + private final int[] partitionCounts; + + /** + * Per topic, its number of subscribers. + */ + private final int[] subscriberCounts; + + /** + * Per topic, the cohorts subscribing to it, in cohort order: an array of cohorts, as long as + * the number of such cohorts. + */ + private final int[][] cohortsOfTopic; + + /** + * Per cohort, its members, in member order: an array of members, as long as the cohort. + */ + private final int[][] cohortMembers; + + /** + * Per cohort, its subscribed topics, in topic order: an array of topics, as long as the + * subscription. + */ + private final int[][] cohortTopics; + + /** + * Reads the members and their subscriptions. + * + * @param groupSpec The group spec. + * @param subscribedTopicDescriber The describer of the subscribed topics. + */ + IndexedGroup(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) { + memberIds = groupSpec.memberIds().toArray(new String[0]); + Arrays.sort(memberIds); + cohortOfMember = new int[memberIds.length]; + + var reader = new SubscriptionReader(groupSpec, subscribedTopicDescriber); + List<int[]> cohortSubscriptions = groupSpec.subscriptionType() == SubscriptionType.HOMOGENEOUS || memberIds.length <= 1 + ? readHomogeneousCohorts(reader) + : readHeterogeneousCohorts(reader); + + // The reader indexes the topics in the order seen, as the topics are only known once every + // subscription is read. They are indexed in topic id order instead, so that the indices + // only depend on the content of the input. The map of the reader is renumbered in place, + // its topics that do not exist keeping -1. + topicIds = reader.topicIds.toArray(new Uuid[0]); + Arrays.sort(topicIds); + topicIndices = reader.topicIndices; + partitionCounts = new int[topicIds.length]; + // Per index in the order seen, the index in topic id order. + var renumbered = new int[topicIds.length]; + for (int topic = 0; topic < topicIds.length; topic++) { + int seenIndex = topicIndices.put(topicIds[topic], topic); + renumbered[seenIndex] = topic; + partitionCounts[topic] = reader.partitionCounts.get(seenIndex); + } + + // The topics of the cohorts are renumbered too, which is cheaper than reading the + // subscriptions again, and the subscribers of every topic are counted. + cohortMembers = membersByCohort(cohortSubscriptions.size()); + cohortTopics = new int[cohortSubscriptions.size()][]; + subscriberCounts = new int[topicIds.length]; + for (int cohort = 0; cohort < cohortSubscriptions.size(); cohort++) { + var topics = cohortSubscriptions.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(); + } + + /** + * Reads the subscription of the first member, which every member has. Review Comment: Reads weirdly ```suggestion * Reads the subscription of the first member, which every member is assumed to share. ``` ########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/GroupModel.java: ########## @@ -0,0 +1,429 @@ +/* + * 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 { + /** + * Reads the subscriptions of the members, numbering their topics that exist in the order they + * are met, and remembering their partition counts. + */ + private static final class SubscriptionReader { + private final GroupSpec groupSpec; + private final SubscribedTopicDescriber describer; + + /** + * Per topic met, its number in the order met, or -1 if it does not exist, so that the + * describer is asked once per topic. + */ + private final Map<Uuid, Integer> indices = new HashMap<>(); + + /** + * The topics met that exist, in the order met. + */ + private final List<Uuid> ids = new ArrayList<>(); + + /** + * The partition counts of the topics met that exist, in the order met. + */ + private final IntList partitionCounts = new IntList(16); + + SubscriptionReader(GroupSpec groupSpec, SubscribedTopicDescriber describer) { + this.groupSpec = groupSpec; + this.describer = describer; + } + + /** + * Reads the subscription of the member. + * + * @param memberId The member id. + * @param topics Receives the numbers of its subscribed topics that exist. + */ + void subscribedTopics(String memberId, IntList topics) { + for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) { + var index = indices.get(topicId); + if (index == null) { + // Met for the first time: numbered next, unless it does not exist. + int partitionCount = describer.numPartitions(topicId); + index = partitionCount < 0 ? -1 : ids.size(); + if (index >= 0) { + ids.add(topicId); + partitionCounts.add(partitionCount); + } + indices.put(topicId, index); + } + if (index >= 0) { + topics.add(index); + } + } + } + } + + /** + * The cohorts found so far, looked up by the topics of their subscription. + * + * <p>A subscription is summarized by a hash that does not depend on the order of its topics, + * and the cohorts with the same hash form a chain. A subscription equals a cohort of its + * chain when it has as many topics and its member subscribes to all the topics of the cohort, + * which the last subscriber of every topic tells without sorting the subscription nor copying + * it. + */ + private static final class CohortsByTopics { Review Comment: Oh I wasn't pushing to revert to a `HashMap<Set, int>`, I just wanted it to be clear that it was more or less equivalent. I was also fine with the manual linked list implementation. But maybe an `IntList` is more readable without sacrificing performance. -- 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]
