squah-confluent commented on code in PR #23732: URL: https://github.com/apache/kafka/pull/23732#discussion_r4254966885
########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/uniform2/IndexedGroup.java: ########## @@ -0,0 +1,525 @@ +/* + * 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 store and process + * every distinct 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; + * an empty group is a single cohort without members. 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 members + * are never assigned to the 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. + * + * <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. + */ + static final class CohortsBySubscription { + /** + * The mask applied to every hash: all the bits in the assignor, none in a test that makes + * every subscription collide. + */ + private final long hashMask; + + /** + * 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; + + CohortsBySubscription() { + this(-1L); + } + + /** + * @param hashMask The mask applied to every hash. + */ + CohortsBySubscription(long hashMask) { + this.hashMask = hashMask; + } + + /** + * @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 = subscriptionHash(subscribedTopics) & hashMask; + 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 arrays are handed over: no subscription + * can be looked up once they are changed. + */ + List<int[]> subscriptions() { + return subscriptions; + } Review Comment: The ownership transfer plus invalidating the `CohortsBySubscription` instance is pretty unexpected for a getter-style method. Is there a way we can make the API less surprising somehow? Things we could try: * Setting an internal "closed" flag that disallows further `getOrCreateCohort` calls * and rename the method to "take(Cohort)Subscriptions" * or turn it into a builder-style API. -- 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]
