gabriellefu commented on code in PR #23790: URL: https://github.com/apache/kafka/pull/23790#discussion_r4297535168
########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/RackAwareStandbyPicker.java: ########## @@ -0,0 +1,445 @@ +/* + * 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.Utils; + +import java.util.ArrayList; +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.TreeSet; +import java.util.function.Function; +import java.util.function.Predicate; +import java.util.function.ToDoubleFunction; +import java.util.stream.IntStream; + +/** + * Picks the conditions for the next standby of one task: the most tag keys that a process with room can meet, + * preferring the keys earlier in {@code rack.aware.assignment.tags}. A process meets a tag key when it has a value + * for it that no holder of the task has yet. It can then tell whether a process meets the conditions and return the + * least-loaded process with room among those that do. The caller reports every change to a load or room with + * {@link #updated(Object)} or {@link #updatedAll()}. + * <p> + * The processes sit in a tree with one level per tag key, those with the same values sharing a leaf. Scanning every + * leaf per standby is too slow once there are many, e.g. one per process for a host tag. Each node keeps its + * least-loaded process with room, and its children are ordered by theirs. A query skips a subtree whose own value is + * used and stops at the first child no lighter than the best process found. + * + * @param <P> The assignor's process type. + */ +final class RackAwareStandbyPicker<P> { + + // The value index of a process for a key it has no value for. + private static final int MISSING_VALUE_INDEX = -1; + // The key index and value index of the root, which has no own value. + private static final int ROOT_PLACEHOLDER_INDEX = -2; + + private final ToDoubleFunction<P> processLoad; + private final Predicate<P> processHasRoom; + // 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; + // Per process, its current entry, which knows its leaf. + private final Map<P, QueuedProcess<P>> queuedByProcess; + 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 startTask. + private final BitSet[] usedValues; + + /** + * @param tagKeys The keys of {@code rack.aware.assignment.tags}; not empty, since without them there is nothing to pick. + * @param processes All processes of the streams group; an earlier one wins a load tie. + * @param clientTags The client tags of a process. + * @param load The load of a process. + * @param hasRoom Whether a process can take another task; {@code process -> true} without a limit. + */ + RackAwareStandbyPicker( + final List<String> tagKeys, + final Collection<P> processes, + final Function<P, Map<String, String>> clientTags, + final ToDoubleFunction<P> load, + final Predicate<P> hasRoom + ) { + if (tagKeys.isEmpty()) { + throw new IllegalArgumentException("Tag keys must not be empty"); + } + processLoad = load; + processHasRoom = hasRoom; + // One level per tag key, so this is also the number of tag keys. + final int treeHeight = tagKeys.size(); + // Per key, the index of each of its values, in the order the values first appear on the processes. + final List<Map<String, Integer>> valueIndexAllKeys = new ArrayList<>(treeHeight); + for (int keyIndex = 0; keyIndex < treeHeight; keyIndex++) { + valueIndexAllKeys.add(new HashMap<>()); + } + // Per process in order, its value index for each key. + final List<int[]> valueIndexesOfProcesses = new ArrayList<>(processes.size()); + for (final P process : processes) { + final Map<String, String> tags = clientTags.apply(process); + final int[] valueIndexes = new int[treeHeight]; + for (int keyIndex = 0; keyIndex < treeHeight; keyIndex++) { + final String value = tags.get(tagKeys.get(keyIndex)); + final Map<String, Integer> valueIndexPerKey = valueIndexAllKeys.get(keyIndex); + valueIndexes[keyIndex] = value == null + ? MISSING_VALUE_INDEX + : valueIndexPerKey.computeIfAbsent(value, __ -> valueIndexPerKey.size()); + } + valueIndexesOfProcesses.add(valueIndexes); + } + + // 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] = valueIndexAllKeys.get(keyIndex).size(); + usedValues[keyIndex] = new BitSet(numValues[keyIndex]); + } + tagKeysInTreeOrder = tagKeysInTreeOrder(numValues); + + root = new Node<>(null, ROOT_PLACEHOLDER_INDEX, ROOT_PLACEHOLDER_INDEX, false, null); + queuedByProcess = Utils.newHashMap(processes.size()); + // The children of each inner node by their value index, only needed to build the tree. + final Map<Node<P>, Map<Integer, Node<P>>> childrenByValueIndex = new HashMap<>(); + int processIndex = 0; + for (final P process : processes) { + insert( + process, + processIndex, + valueIndexesOfProcesses.get(processIndex), + childrenByValueIndex + ); + processIndex++; + } + } + + /** Adds the process to the leaf of its values, creating the leaf and the inner nodes on its path. */ + private void insert( + final P process, + final int processIndex, + final int[] valueIndexes, + final Map<Node<P>, Map<Integer, Node<P>>> childrenByValueIndex + ) { + 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]; + final Map<Integer, Node<P>> children = childrenByValueIndex.computeIfAbsent(node, __ -> new HashMap<>()); + Node<P> child = children.get(valueIndex); + if (child == null) { + child = new Node<>(node, keyIndex, valueIndex, level == treeHeight - 1, valueIndexes); + children.put(valueIndex, child); + } + node = child; + } + final QueuedProcess<P> queued = new QueuedProcess<>(process, processIndex, node, processLoad.applyAsDouble(process)); + if (processHasRoom.test(process)) { + node.processesByLoad.add(queued); + repairPath(node); + } + queuedByProcess.put(process, queued); + } + + /** 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) { + return IntStream.range(0, numValues.length) + .boxed() + .sorted(Comparator.comparingInt(keyIndex -> numValues[keyIndex])) + .mapToInt(Integer::intValue) + .toArray(); + } + + /** The tag keys in priority order: that of {@code rack.aware.assignment.tags}, not of the levels. */ + List<TagKey> tagKeys() { + return keys; + } + + /** + * Starts the picks for a task, forgetting the values of the holders of the previous task; a missing value stays used. + * + * @param holders The holders of the task: its active owner and the standbys placed so far. + */ + void startTask(final Collection<P> holders) { + for (final BitSet used : usedValues) { + used.clear(); + } + for (final P holder : holders) { + markUsed(holder); + } + } + + /** Records the tag values of a new holder of the task as used. */ + void markUsed(final P holder) { + final int[] valueIndexes = queuedByProcess.get(holder).leaf.valueIndexes; + for (int keyIndex = 0; keyIndex < keys.size(); keyIndex++) { + // A missing value has no bit: isUsed counts it as used. + if (valueIndexes[keyIndex] != MISSING_VALUE_INDEX) { + usedValues[keyIndex].set(valueIndexes[keyIndex]); + } + } + } + + /** Reads the load and room of the process again; call it after either changes. */ + void updated(final P process) { + queuedByProcess.put(process, requeue(queuedByProcess.get(process))); + } + + /** Reads the load and room of every process again; call it after a change that may affect many. */ + void updatedAll() { + queuedByProcess.replaceAll((process, queued) -> requeue(queued)); + } + + /** The conditions for the next standby, or empty once no process with room can make the task more diverse. */ + List<TagKey> pick() { + // No process has room for the standby. + if (root.leastLoaded == null) { + return List.of(); + } + // Only if a tag key has an unused value is it possible to choose a process that makes the task more diverse + // in this tag key. + final List<TagKey> tagKeysWithUnusedValue = new ArrayList<>(); + for (final TagKey tagKey : keys) { + if (hasUnusedValue(tagKey)) { + tagKeysWithUnusedValue.add(tagKey); + } + } + if (tagKeysWithUnusedValue.isEmpty()) { + return List.of(); + } + + // Most picks find a process with room that meets the condition of every such tag key. + if (leastLoadedProcessMeeting(tagKeysWithUnusedValue) != null) { Review Comment: I have change pick to return a record record Pick<P>(List<TagKey> conditions, P leastLoaded) which return both the condition and the leastLoaded process so there's no duplicate search -- 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]
