gabriellefu commented on code in PR #23790: URL: https://github.com/apache/kafka/pull/23790#discussion_r4262581941
########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/TagTree.java: ########## @@ -0,0 +1,423 @@ +/* + * 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.streams.assignor; + +import org.apache.kafka.coordinator.group.streams.assignor.IdenticalTagGroups.QueuedProcess; +import org.apache.kafka.coordinator.group.streams.assignor.IdenticalTagGroups.TagGroup; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.BitSet; +import java.util.Collection; +import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.PriorityQueue; + +/** + * A tree built from {@code rack.aware.assignment.tags} with every tag group of {@link IdenticalTagGroups} as a leaf. + * Given some tag keys as conditions, it tells whether a tag group with room has values for them that no holder of the + * task has yet, and returns the least-loaded process with room among such tag groups. It assumes that loads only grow + * and room only shrinks. + * + * @param <P> The assignor's process type. + */ +final class TagTree<P> { + + // The keys in the order of rack.aware.assignment.tags. + private final List<TagKey> keys; + // Per key index, the number of values. + private final int[] numValues; + // Per level from the top, the index of its key. The priority of a key is still its order in + // rack.aware.assignment.tags, not this one. + private final int[] tagKeysInTreeOrder; + private final Map<TagGroup<P>, Node<P>> leafByTagGroup; + private final Node<P> root; + + // Per key index, one bit per value index, set once a holder of the task carries that value. + // Cleared by clearUsedValues. + private final BitSet[] usedValues; + // The nodes with a cached least-loaded process, to clear on startPick: loads do not change during a pick. + private final List<Node<P>> cachedNodes = new ArrayList<>(); + + /** + * @param tagKeys The keys of {@code rack.aware.assignment.tags}. + * @param tagGroups All tag groups of the streams group. + */ + TagTree(final List<String> tagKeys, final Collection<TagGroup<P>> tagGroups) { + // One level per tag key, so this is also the number of tag keys. + final int treeHeight = tagKeys.size(); + final List<Map<String, Integer>> indexOfValues = indexValues(tagKeys, tagGroups); + + // The keys and the number of values of each, with no value used yet. + keys = new ArrayList<>(treeHeight); + numValues = new int[treeHeight]; + usedValues = new BitSet[treeHeight]; + for (int keyIndex = 0; keyIndex < treeHeight; keyIndex++) { + keys.add(new TagKey(tagKeys.get(keyIndex), keyIndex)); + numValues[keyIndex] = indexOfValues.get(keyIndex).size(); + usedValues[keyIndex] = new BitSet(numValues[keyIndex] + 1); + } + tagKeysInTreeOrder = tagKeysInTreeOrder(numValues); + + root = new Node<>(-1, -1, treeHeight, false); + leafByTagGroup = new HashMap<>(); + for (final TagGroup<P> tagGroup : tagGroups) { + insert(tagGroup, valueIndexes(tagGroup, tagKeys, indexOfValues)); + } + buildInnerNodes(root); + } + + /** Per key, the index of each of its values, in the order the values first appear on the tag groups. */ + private static <P> List<Map<String, Integer>> indexValues(final List<String> tagKeys, final Collection<TagGroup<P>> tagGroups) { + final List<Map<String, Integer>> indexOfValues = new ArrayList<>(tagKeys.size()); + for (final String tagKey : tagKeys) { + final Map<String, Integer> indexOfValue = new HashMap<>(); + for (final TagGroup<P> tagGroup : tagGroups) { + final String value = tagGroup.clientTags().get(tagKey); + if (value != null) { + indexOfValue.putIfAbsent(value, indexOfValue.size()); + } + } + indexOfValues.add(indexOfValue); + } + return indexOfValues; + } + + /** The value index of the tag group for each key, the missing value's index for a key it has no value for. */ + private int[] valueIndexes(final TagGroup<P> tagGroup, final List<String> tagKeys, final List<Map<String, Integer>> indexOfValues) { + final int[] valueIndexes = new int[tagKeys.size()]; + for (int keyIndex = 0; keyIndex < valueIndexes.length; keyIndex++) { + final String value = tagGroup.clientTags().get(tagKeys.get(keyIndex)); + valueIndexes[keyIndex] = indexOfValues.get(keyIndex).getOrDefault(value, missingValueIndex(keyIndex)); + } + return valueIndexes; + } + + /** Adds the tag group as a leaf, creating the inner nodes on its path. */ + private void insert(final TagGroup<P> tagGroup, final int[] valueIndexes) { + final int treeHeight = tagKeysInTreeOrder.length; + Node<P> node = root; + for (int level = 0; level < treeHeight; level++) { + final int keyIndex = tagKeysInTreeOrder[level]; + final int valueIndex = valueIndexes[keyIndex]; + Node<P> child = node.childrenByValueIndex.get(valueIndex); + if (child == null) { + child = new Node<>(keyIndex, valueIndex, treeHeight, level == treeHeight - 1); + node.childrenByValueIndex.put(valueIndex, child); + } + node = child; + } + node.tagGroup = tagGroup; + node.valueIndexes = valueIndexes; + leafByTagGroup.put(tagGroup, node); + } + + /** The key indexes as the levels from the top, fewest values first; the order affects speed, not results. */ + private static int[] tagKeysInTreeOrder(final int[] numValues) { + final Integer[] keyIndexes = new Integer[numValues.length]; + for (int keyIndex = 0; keyIndex < keyIndexes.length; keyIndex++) { + keyIndexes[keyIndex] = keyIndex; + } + Arrays.sort( + keyIndexes, + Comparator.<Integer>comparingInt(keyIndex -> numValues[keyIndex]).thenComparingInt(keyIndex -> keyIndex) + ); + final int[] tagKeysInTreeOrder = new int[keyIndexes.length]; + for (int level = 0; level < keyIndexes.length; level++) { + tagKeysInTreeOrder[level] = keyIndexes[level]; + } + return tagKeysInTreeOrder; + } + + /** Fills in the values below and the child queues of the node and of every inner node under it. */ + private void buildInnerNodes(final Node<P> node) { + if (node.tagGroup != null) { + return; + } + for (final Node<P> child : node.childrenByValueIndex.values()) { + buildInnerNodes(child); + // The child's own value, then the values below it. + node.valuesBelowOf(child.keyIndex).set(child.valueIndex); + for (int keyIndex = 0; keyIndex < child.valuesBelow.length; keyIndex++) { + if (child.valuesBelow[keyIndex] != null) { + node.valuesBelowOf(keyIndex).or(child.valuesBelow[keyIndex]); + } + } + // A child without room below never gets room again, so it is not queued. + final QueuedProcess<P> childLeastLoaded = leastLoadedProcessBelow(child); + if (childLeastLoaded != null) { + child.queuedLoad = childLeastLoaded.load; + child.queuedProcessIndex = childLeastLoaded.processIndex; + node.childrenByLeastLoaded.add(child); + } + } + } + + /** The tag keys in priority order: that of {@code rack.aware.assignment.tags}, not of the levels. */ + List<TagKey> tagKeys() { + return keys; + } + + /** Forgets the values of the holders of the previous task; a missing value stays used. */ + void clearUsedValues() { + for (final BitSet used : usedValues) { + used.clear(); + } + } + + /** Records the values of a holder of the task as used. */ + void markUsed(final TagGroup<P> holder) { + final int[] valueIndexes = leafByTagGroup.get(holder).valueIndexes; + for (int keyIndex = 0; keyIndex < keys.size(); keyIndex++) { + usedValues[keyIndex].set(valueIndexes[keyIndex]); + } + } + + /** Whether some value of the key is not used yet. */ + boolean hasUnusedValue(final TagKey key) { + // The first unused value index is a value, not the missing one after the last value. + return usedValues[key.index].nextClearBit(0) < numValues[key.index]; + } + + /** Whether the values of the tag group for {@code conditionKeys} are all unused. */ + boolean tagGroupMeeting(final TagGroup<P> tagGroup, final List<TagKey> conditionKeys) { + final int[] valueIndexes = leafByTagGroup.get(tagGroup).valueIndexes; + for (final TagKey key : conditionKeys) { + if (valueUsed(key.index, valueIndexes[key.index])) { + return false; + } + } + return true; + } + + /** The value index of a missing value, after the last value. */ + private int missingValueIndex(final int keyIndex) { + return numValues[keyIndex]; + } + + /** Whether the value is used: a holder of the task carries it, or it is missing. */ + private boolean valueUsed(final int keyIndex, final int valueIndex) { + return valueIndex == missingValueIndex(keyIndex) || usedValues[keyIndex].get(valueIndex); + } + + /** Clears the cache, since loads and room may have changed since the last pick. */ + void startPick() { + for (final Node<P> node : cachedNodes) { + node.cached = false; + } + cachedNodes.clear(); + } + + /** Whether a tag group with room has values for {@code conditionKeys} that are all unused. */ + boolean hasTagGroupMeeting(final List<TagKey> conditionKeys) { + return hasTagGroupBelowMeeting(root, keyBits(conditionKeys)); + } + + /** The least-loaded process with room in a tag group meeting {@code conditionKeys}, or null if none. */ + QueuedProcess<P> leastLoadedProcessMeeting(final List<TagKey> conditionKeys) { + return leastLoadedProcessBelowMeeting(root, keyBits(conditionKeys)); + } + + /** The keys as an int with bit k set for the key with index k. */ + private static int keyBits(final List<TagKey> keys) { + int keyBits = 0; + for (final TagKey key : keys) { + keyBits |= 1 << key.index; Review Comment: Thanks.a Kafka Streams client can set at most 5 client tags (StreamsConfig.MAX_RACK_AWARE_ASSIGNMENT_TAG_LIST_SIZE). Capping the broker-side config would change a config that is already in 4.4, so I'm leaving it out of this PR. -- 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]
