lucasbru commented on code in PR #23790:
URL: https://github.com/apache/kafka/pull/23790#discussion_r4256691652


##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/RackAwareStandbyPicker.java:
##########
@@ -0,0 +1,101 @@
+/*
+ * 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.TagGroup;
+import org.apache.kafka.coordinator.group.streams.assignor.TagTree.TagKey;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+/**
+ * Picks the conditions for the next standby of one task: the most tag keys 
that a tag group with room can meet,
+ * preferring the keys earlier in {@code rack.aware.assignment.tags}. A tag 
group 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 tag 
group meets the conditions and return the
+ * least-loaded process with room among those that do.
+ *
+ * @param <P> The assignor's process type.
+ */
+final class RackAwareStandbyPicker<P> {
+
+    private final TagTree<P> tree;
+
+    /**
+     * @param tree    The tag tree of the assignment, whose used values the 
picker resets for this task.
+     * @param holders The holders of the task: its active owner and the 
standbys placed so far.
+     */
+    RackAwareStandbyPicker(final TagTree<P> tree, final 
Collection<TagGroup<P>> holders) {
+        this.tree = tree;
+        tree.clearUsedValues();
+        for (final TagGroup<P> holder : holders) {
+            markUsed(holder);
+        }
+    }
+
+    /** Records the tag values of a new holder of the task as used. */
+    void markUsed(final TagGroup<P> holder) {
+        tree.markUsed(holder);
+    }
+
+    /** The conditions for the next standby, or empty once no tag group with 
room can make the task more diverse. */
+    List<TagKey> pick() {
+        // 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 : tree.tagKeys()) {
+            if (tree.hasUnusedValue(tagKey)) {
+                tagKeysWithUnusedValue.add(tagKey);
+            }
+        }
+        if (tagKeysWithUnusedValue.isEmpty()) {
+            return List.of();
+        }
+        tree.startPick();
+        // No tag group has room for the standby; with no conditions, only 
room counts.
+        if (!tree.hasTagGroupMeeting(List.of())) {
+            return List.of();
+        }
+
+        // Most picks find a tag group with room that meets the condition of 
every such tag key.
+        List<TagKey> conditions = tagKeysWithUnusedValue;
+        if (!tree.hasTagGroupMeeting(conditions)) {
+            // Otherwise the tag keys are added in priority order; a tag key 
is given up when no tag group with room
+            // meets it together with the tag keys kept so far.
+            conditions = new ArrayList<>();
+            for (final TagKey tagKey : tagKeysWithUnusedValue) {
+                final List<TagKey> withTagKey = new ArrayList<>(conditions);
+                withTagKey.add(tagKey);
+                // If withTagKey holds every tag key, the query above already 
found no tag group for it.
+                if (withTagKey.size() < tagKeysWithUnusedValue.size() && 
tree.hasTagGroupMeeting(withTagKey)) {
+                    conditions = withTagKey;
+                }
+            }
+        }
+        return conditions;
+    }
+
+    /** Whether a tag group is a candidate under the conditions of a pick: it 
meets them and has room. */
+    boolean isCandidate(final TagGroup<P> tagGroup, final List<TagKey> 
conditions) {
+        return tree.tagGroupMeeting(tagGroup, conditions) && 
tagGroup.hasRoom();
+    }
+
+    /** The least-loaded process with room of the candidates; call it before 
placing the standby. */
+    P leastLoaded(final List<TagKey> conditions) {
+        return tree.leastLoadedProcessMeeting(conditions).process;

Review Comment:
   `pick()` returns an empty list both when no tag group has room and when no 
group would add diversity. In the first case `leastLoaded(List.of())` throws an 
NPE here (`leastLoadedProcessMeeting` returns null). In the second it returns 
any process with room, which can be the active owner itself. Could we document 
that `leastLoaded` is only valid for the conditions `pick()` returned, or have 
it return null like `leastLoadedProcessMeeting` does?



##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/IdenticalTagGroups.java:
##########
@@ -0,0 +1,142 @@
+/*
+ * 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.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.function.Function;
+import java.util.function.Predicate;
+import java.util.function.ToDoubleFunction;
+
+/**
+ * Splits the processes of a streams group into tag groups by their values for 
{@code rack.aware.assignment.tags}.
+ * The least-loaded lookups assume that loads only grow and room only shrinks.
+ *
+ * @param <P> The assignor's process type.
+ */
+final class IdenticalTagGroups<P> {
+
+    private final Collection<TagGroup<P>> tagGroups;
+    private final Map<P, TagGroup<P>> tagGroupByProcess;
+
+    /**
+     * @param tagKeys    The keys of {@code rack.aware.assignment.tags}.
+     * @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.
+     */
+    IdenticalTagGroups(
+        final List<String> tagKeys,
+        final Collection<P> processes,
+        final Function<P, Map<String, String>> clientTags,
+        final ToDoubleFunction<P> load,
+        final Predicate<P> hasRoom
+    ) {
+        final Map<List<String>, TagGroup<P>> tagGroupsByTagValues = new 
HashMap<>();
+        tagGroupByProcess = Utils.newHashMap(processes.size());
+        int processIndex = 0;
+        for (final P process : processes) {
+            final Map<String, String> tags = clientTags.apply(process);
+            final List<String> tagValues = new ArrayList<>(tagKeys.size());
+            for (final String tagKey : tagKeys) {
+                tagValues.add(tags.get(tagKey));
+            }
+            final TagGroup<P> tagGroup =
+                tagGroupsByTagValues.computeIfAbsent(tagValues, values -> new 
TagGroup<>(tags, load, hasRoom));
+            tagGroup.processesByLoad.add(new QueuedProcess<>(process, 
processIndex++, load.applyAsDouble(process)));
+            tagGroupByProcess.put(process, tagGroup);
+        }
+        tagGroups = tagGroupsByTagValues.values();
+    }
+
+    Collection<TagGroup<P>> tagGroups() {
+        return tagGroups;
+    }
+
+    /** The tag group of a process, valid even when loads drop. */
+    TagGroup<P> tagGroupOf(final P process) {
+        return tagGroupByProcess.get(process);
+    }
+
+    /** Processes with the same values for {@code rack.aware.assignment.tags}. 
*/
+    static final class TagGroup<P> {
+        private final Map<String, String> clientTags;
+        private final ToDoubleFunction<P> processLoad;
+        private final Predicate<P> processHasRoom;
+        // The processes that may still have room, by the load each was queued 
with.
+        private final PriorityQueue<QueuedProcess<P>> processesByLoad = new 
PriorityQueue<>(QueuedProcess.ORDER);
+
+        private TagGroup(final Map<String, String> clientTags, final 
ToDoubleFunction<P> processLoad, final Predicate<P> processHasRoom) {
+            this.clientTags = clientTags;
+            this.processLoad = processLoad;
+            this.processHasRoom = processHasRoom;
+        }
+
+        /** The client tags of the first process: only the values for {@code 
rack.aware.assignment.tags} are shared by the tag group. */
+        Map<String, String> clientTags() {

Review Comment:
   `clientTags()` exposes the whole tag map of the first process, with a caveat 
that only some entries are shared, and `TagTree` then re-reads it with its own 
`tagKeys` list. The constructor already computes `tagValues` (the map key at 
line 67). Could `TagGroup` keep that `List<String>` instead, aligned with the 
tag keys? It would be harder to misuse and the two classes wouldn't depend on 
being handed the same key list.



##########
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;
+        }
+        return keyBits;
+    }
+
+    /** Whether a tag group with room below the node meets {@code 
conditionKeys}. */
+    private boolean hasTagGroupBelowMeeting(final Node<P> node, final int 
conditionKeys) {
+        if (valuesBelowAllMeeting(node, conditionKeys)) {
+            return leastLoadedProcessBelow(node) != null;
+        }
+        // Walk into the children in queue order until one has a tag group 
meeting the conditions.
+        boolean found = false;
+        while (!found && nextChildMeeting(node, conditionKeys, null) != null) {
+            found = hasTagGroupBelowMeeting(setAsideHead(node), conditionKeys);
+        }
+        putBackSetAsideChildren(node);
+        return found;
+    }
+
+    /**
+     * The least-loaded process with room below the node in a tag group 
meeting {@code conditionKeys}, or null if none.
+     */
+    private QueuedProcess<P> leastLoadedProcessBelowMeeting(final Node<P> 
node, final int conditionKeys) {
+        if (valuesBelowAllMeeting(node, conditionKeys)) {
+            return leastLoadedProcessBelow(node);
+        }
+        // Walk into the children in queue order until the next one cannot 
beat the best found.
+        QueuedProcess<P> leastLoaded = null;
+        while (nextChildMeeting(node, conditionKeys, leastLoaded) != null) {
+            final QueuedProcess<P> found = 
leastLoadedProcessBelowMeeting(setAsideHead(node), conditionKeys);
+            if (found != null && (leastLoaded == null || 
QueuedProcess.ORDER.compare(found, leastLoaded) < 0)) {
+                leastLoaded = found;
+            }
+        }
+        putBackSetAsideChildren(node);
+        return leastLoaded;
+    }
+
+    /** Takes the child at the head of the node's queue out of it until {@link 
#putBackSetAsideChildren(Node)}. */
+    private Node<P> setAsideHead(final Node<P> node) {
+        final Node<P> head = node.childrenByLeastLoaded.poll();
+        node.setAsideChildren.add(head);
+        return head;
+    }
+
+    /** Puts the children that a query set aside back in the node's queue. */
+    private void putBackSetAsideChildren(final Node<P> node) {
+        node.childrenByLeastLoaded.addAll(node.setAsideChildren);
+        node.setAsideChildren.clear();
+    }
+
+    /** Whether every value below the node meets {@code conditionKeys}. */
+    private boolean valuesBelowAllMeeting(final Node<P> node, final int 
conditionKeys) {
+        // One key of conditionKeys at a time, lowest bit first.
+        for (int keyBits = conditionKeys; keyBits != 0; keyBits &= keyBits - 
1) {
+            final int keyIndex = Integer.numberOfTrailingZeros(keyBits);
+            final BitSet valuesBelow = node.valuesBelow[keyIndex];
+            if (valuesBelow == null) {
+                continue;
+            }
+            if (valuesBelow.get(missingValueIndex(keyIndex)) || 
valuesBelow.intersects(usedValues[keyIndex])) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    /** The least-loaded process with room below the node, or null if none; 
computed once per pick. */
+    private QueuedProcess<P> leastLoadedProcessBelow(final Node<P> node) {
+        if (!node.cached) {
+            if (node.tagGroup != null) {
+                node.cachedLeastLoaded = node.tagGroup.leastLoadedWithRoom();
+            } else {
+                // With no conditions, the next child is the one with the 
least-loaded process below it.
+                final Node<P> head = nextChildMeeting(node, 0, null);
+                node.cachedLeastLoaded = head == null ? null : 
leastLoadedProcessBelow(head);
+            }
+            node.cached = true;
+            cachedNodes.add(node);
+        }
+        return node.cachedLeastLoaded;
+    }
+
+    /**
+     * Brings the next child whose own value meets {@code conditionKeys} to 
the head of the queue and returns it, or
+     * null if none left can beat {@code bound}.
+     */
+    private Node<P> nextChildMeeting(final Node<P> node, final int 
conditionKeys, final QueuedProcess<P> bound) {
+        while (!node.childrenByLeastLoaded.isEmpty()) {
+            final Node<P> child = node.childrenByLeastLoaded.peek();
+            // The process a child is queued with may be stale, but a stale 
one is only lighter than its current one.
+            if (bound != null) {
+                // Lighter by QueuedProcess.ORDER: by load, then by process 
index.
+                final int byLoad = Double.compare(child.queuedLoad, 
bound.load);

Review Comment:
   nit: the (load, processIndex) ordering is written out three times: 
`QueuedProcess.ORDER`, `Node.BY_LEAST_LOADED` and here. Could one static 
compare helper be shared?



##########
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> {

Review Comment:
   This is a fair amount of machinery (tree, per-node queues, set-aside lists, 
a per-pick cache), and the test's oracle is a linear scan. Do you have a sense 
of the group sizes where a scan over tag groups gets too slow? A short note in 
the class Javadoc on why the tree is needed would help future readers.



-- 
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]

Reply via email to