KAFKA-19909: Add rack awareness assignment to UniformHomogeneousAssignmentBuilder

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignorHelpers.java

@@ -17,11 +17,13 @@
 package org.apache.kafka.coordinator.group.assignor;
 
 import org.apache.kafka.common.Uuid;
+import org.apache.kafka.server.common.TopicIdPartition;
 
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 
 /**
@@ -69,4 +71,34 @@ static <K, V> HashMap<K, V> newHashMap(int numMappings) {
     static <K> HashSet<K> newHashSet(int numElements) {
         return new HashSet<>((int) (((numElements + 1) / 0.75f) + 1));
     }
+
+
+    /**
+     * Checks if the member's rack matches any of the partition's racks.
+     * @param memberRackId The member's rack id.
+     * @param partitionRackIds The partition's rack ids.
+     * @return True if the member's rack matches any of the partition's racks, false otherwise.
+     */
+    public static boolean isRackMatch(Optional<String> memberRackId, Set<String> partitionRackIds) {
+        return memberRackId.isPresent() && partitionRackIds.contains(memberRackId.get());
+    }
+
+    /**
+     * Determines whether rack-aware assignment should be used based on the provided racks.
+     * @param allMemberRacks The set of all member racks.
+     * @param allPartitionRacks The set of all partition racks.
+     * @param racksPerPartition A map of partitions to their respective racks.
+     * @return True if member racks and partition racks overlap and not all partitions have the same set of racks, false otherwise.
+     */
+    public static boolean useRackAwareAssignment(
+        Set<String> allMemberRacks,
+        Set<String> allPartitionRacks,
+        Map<TopicIdPartition, Set<String>> racksPerPartition
+    ) {
+        if (allMemberRacks.isEmpty() || Collections.disjoint(allMemberRacks, allPartitionRacks))
+            return false;
+        else {
+            return !racksPerPartition.values().stream().allMatch(allPartitionRacks::equals);
+        }
+    }
 }

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformHomogeneousAssignmentBuilder.java

@@ -41,11 +41,13 @@
  * <li> Balance:          Ensure partitions are distributed equally among all members.
  *                        The difference in assignments sizes between any two members
  *                        should not exceed one partition. </li>
+ * <li> Rack Awareness:   When feasible, aim to assign partitions to members
+ *                        located on the same rack thus avoiding cross-zone traffic. </li>
  * <li> Stickiness:       Minimize partition movements among members by retaining
  *                        as much of the existing assignment as possible. </li>
  *
  * The assignment builder prioritizes the properties in the following order:
- *      Balance > Stickiness.
+ *      Balance > Rack Awareness > Stickiness.
  */
 public class UniformHomogeneousAssignmentBuilder {
     /**
@@ -72,7 +74,7 @@ public class UniformHomogeneousAssignmentBuilder {
      * The partitions that still need to be assigned.
      * Initially this contains all the subscribed topics' partitions.
      */
-    private final List<TopicIdPartition> unassignedPartitions;
+    private List<TopicIdPartition> unassignedPartitions;
 
     /**
      * The target assignment.
@@ -85,6 +87,16 @@ public class UniformHomogeneousAssignmentBuilder {
      */
     private int minimumMemberQuota;
 
+    /**
+     * Whether to use rack aware assignment strategy.
+     */
+    private final boolean useRackStrategy;
+
+    /**
+     * The mapping of topic partitions to their associated racks.
+     */
+    private final Map<TopicIdPartition, Set<String>> partitionRacks;
+
     /**
      * The number of members to receive an extra partition beyond the minimum quota.
      * Example: If there are 11 partitions to be distributed among 3 members,
@@ -101,6 +113,24 @@ public class UniformHomogeneousAssignmentBuilder {
         this.unassignedPartitions = new ArrayList<>();
 
         this.targetAssignment = new HashMap<>();
+        this.partitionRacks = new HashMap<>();
+
+        Set<String> allMemberRacks = new HashSet<>();
+        for (String memberId : groupSpec.memberIds()) {
+            groupSpec.memberSubscription(memberId).rackId().ifPresent(allMemberRacks::add);
+        }
+
+        Set<String> allPartitionRacks = new HashSet<>();
+        for (Uuid topicId : this.subscribedTopicIds) {
+            int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+            for (int partitionId = 0; partitionId < partitionCount; partitionId++) {
+                Set<String> racks = subscribedTopicDescriber.racksForPartition(topicId, partitionId);
+                partitionRacks.put(new TopicIdPartition(topicId, partitionId), racks);
+                allPartitionRacks.addAll(racks);
+            }
+        }
+
+        this.useRackStrategy = AssignorHelpers.useRackAwareAssignment(allMemberRacks, allPartitionRacks, partitionRacks);
     }
 
     /**
@@ -139,6 +169,11 @@ public GroupAssignment build() throws PartitionAssignorException {
         // exceed the maximum quota assigned to each member.
         maybeRevokePartitions();
 
+        // Assign the unassigned partitions to the members if rack matches.
+        if (useRackStrategy) {
+            assignRackAwarenessRemainingPartitions();
+        }
+
         // Assign the unassigned partitions to the members with space.
         assignRemainingPartitions();
 
@@ -154,11 +189,7 @@ public GroupAssignment build() throws PartitionAssignorException {
      */
     @SuppressWarnings({"CyclomaticComplexity", "NPathComplexity"})
     private void maybeRevokePartitions() {
-        int memberCount = groupSpec.memberIds().size();
-        int memberIndex = -1;
         for (String memberId : groupSpec.memberIds()) {
-            memberIndex++;
-
             Map<Uuid, Set<Integer>> oldAssignment = groupSpec.memberAssignment(memberId).partitions();
             Map<Uuid, Set<Integer>> newAssignment = null;
 
@@ -180,11 +211,17 @@ private void maybeRevokePartitions() {
                 Set<Integer> partitions = topicPartitions.getValue();
 
                 if (subscribedTopicIds.contains(topicId)) {
-                    if (partitions.size() <= quota) {
+                    if (partitions.size() <= quota && !useRackStrategy) {
                         quota -= partitions.size();
                     } else {
                         for (Integer partition : partitions) {
-                            if (quota > 0) {
+                            // Keep the partition when:
+                            // 1. We still have quota left to keep it.
+                            // 2-1. Don't use rack strategy, so we can keep it.
+                            // 2-2. Use rack strategy and the member rack matches the partition racks.
+                            if (quota > 0 && (!useRackStrategy || AssignorHelpers.isRackMatch(
+                                groupSpec.memberSubscription(memberId).rackId(),
+                                partitionRacks.getOrDefault(new TopicIdPartition(topicId, partition), Set.of())))) {
                                 quota--;
                             } else {
                                 if (newAssignment == null) {
@@ -215,21 +252,21 @@ private void maybeRevokePartitions() {
             }
 
             if (quota > 0 &&
-                quotaHasExtraPartition &&
-                memberCount - memberIndex > remainingMembersToGetAnExtraPartition) {
-                // Give up the extra partition quota for another member to claim,
-                // unless this member is one of the last remainingMembersToGetAnExtraPartition
-                // members in the list and must take the extra partition.
+                quotaHasExtraPartition) {
+                // Give up the extra partition quota for another member to claim.
                 quota--;
                 quotaHasExtraPartition = false;
             }
 
-            if (quota > 0) {
-                unfilledMembers.add(new MemberWithRemainingQuota(memberId, quota));
+            if (quota == 0 && quotaHasExtraPartition) {
+                remainingMembersToGetAnExtraPartition--;
             }
 
-            if (quotaHasExtraPartition) {
-                remainingMembersToGetAnExtraPartition--;
+            // An unfilled member is a member that can still take more partitions.
+            // 1. It has quota left.
+            // 2. It doesn't have quota left, but there are remaining extra partition, and it hasn't claimed one yet.
+            if (quota > 0 || (remainingMembersToGetAnExtraPartition > 0 && !quotaHasExtraPartition)) {
+                unfilledMembers.add(new MemberWithRemainingQuota(memberId, quota));
             }
 
             if (newAssignment == null) {
@@ -240,13 +277,75 @@ private void maybeRevokePartitions() {
         }
     }
 
+    /**
+     * Assign the unassigned partitions to the unfilled members if member and partition racks are matched.
+     */
+    private void assignRackAwarenessRemainingPartitions() {
+        // Assign partitions to members with descending order. This avoids the cost of shifting elements.
+        for (int i = unassignedPartitions.size() - 1; i >= 0; i--) {
+            TopicIdPartition tip = unassignedPartitions.get(i);
+            boolean isPartitionAssigned = false;
+
+            for (var unfilledMembersIter = unfilledMembers.iterator(); unfilledMembersIter.hasNext(); ) {
+                MemberWithRemainingQuota unfilledMember = unfilledMembersIter.next();
+                if (unfilledMember.remainingQuota() == 0 && remainingMembersToGetAnExtraPartition == 0) {
+                    unfilledMembersIter.remove();
+                    continue;
+                }
+
+                String memberId = unfilledMember.memberId;
+                if (!AssignorHelpers.isRackMatch(groupSpec.memberSubscription(memberId).rackId(),
+                    partitionRacks.getOrDefault(tip, Set.of()))) {
+                    continue;
+                }
+
+                Map<Uuid, Set<Integer>> newAssignment = targetAssignment.get(memberId).partitions();
+                if (AssignorHelpers.isImmutableMap(newAssignment)) {
+                    // If the new assignment is immutable, we must create a deep copy of it
+                    // before altering it.
+                    newAssignment = AssignorHelpers.deepCopyAssignment(newAssignment);
+                    targetAssignment.put(memberId, new MemberAssignmentImpl(newAssignment));
+                }
+                newAssignment
+                    .computeIfAbsent(tip.topicId(), __ -> new HashSet<>())
+                    .add(tip.partitionId());
+
+                if (unfilledMember.remainingQuota() > 0) {
+                    unfilledMember.decrementQuota();
+                } else {
+                    unfilledMembersIter.remove();
+                    remainingMembersToGetAnExtraPartition--;
+                }
+
+                isPartitionAssigned = true;
+                break;
+            }
+
+            if (isPartitionAssigned) {
+                unassignedPartitions.remove(i);
+            }
+        }
+    }
+
     /**
      * Assign the unassigned partitions to the unfilled members.
      */
     private void assignRemainingPartitions() {
         int unassignedPartitionIndex = 0;
 
+        // If member is one of the last remainingMembersToGetAnExtraPartition, it must take the extra partition.
+        for (int i = unfilledMembers.size() - 1; i >= 0; i--) {
+            if (remainingMembersToGetAnExtraPartition == 0) {
+                break;
+            }
+            unfilledMembers.get(i).incrementQuota();
+            remainingMembersToGetAnExtraPartition--;
+        }
+
         for (MemberWithRemainingQuota unfilledMember : unfilledMembers) {
+            if (unfilledMember.remainingQuota() == 0) {
+                continue;
+            }
             String memberId = unfilledMember.memberId;
             int remainingQuota = unfilledMember.remainingQuota;
 
@@ -272,6 +371,25 @@ private void assignRemainingPartitions() {
         }
     }
 
-    private record MemberWithRemainingQuota(String memberId, int remainingQuota) {
+    private static class MemberWithRemainingQuota {
+        private final String memberId;
+        private int remainingQuota;
+
+        MemberWithRemainingQuota(String memberId, int remainingQuota) {
+            this.memberId = memberId;
+            this.remainingQuota = remainingQuota;
+        }
+
+        int remainingQuota() {
+            return remainingQuota;
+        }
+
+        void incrementQuota() {
+            remainingQuota++;
+        }
+
+        void decrementQuota() {
+            remainingQuota--;
+        }
     }
 }

group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/AssignorHelpersTest.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.assignor;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.server.common.TopicIdPartition;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class AssignorHelpersTest {
+    private static final Uuid TOPIC_ID = Uuid.randomUuid();
+
+    @Test
+    void testIsRackMatchWithEmptyMemberRackId() {
+        assertFalse(AssignorHelpers.isRackMatch(Optional.empty(), Set.of("rack1", "rack2")));
+    }
+
+    @Test
+    void testIsRackMatchWithMatchingRack() {
+        assertTrue(AssignorHelpers.isRackMatch(Optional.of("rack1"), Set.of("rack1", "rack2")));
+    }
+
+    @Test
+    void testIsRackMatchWithNonMatchingRack() {
+        assertFalse(AssignorHelpers.isRackMatch(Optional.of("rack3"), Set.of("rack1", "rack2")));
+    }
+
+    @Test
+    void testIsRackMatchWithEmptyPartitionRacks() {
+        assertFalse(AssignorHelpers.isRackMatch(Optional.of("rack1"), Set.of()));
+    }
+
+    @Test
+    void testUseRackAwareAssignmentWithEmptyMemberRacks() {
+        Set<String> allMemberRacks = Set.of();
+        Set<String> allPartitionRacks = Set.of("rack1", "rack2");
+        Map<TopicIdPartition, Set<String>> racksPerPartition = Map.of(
+            new TopicIdPartition(TOPIC_ID, 0), Set.of("rack1"),
+            new TopicIdPartition(TOPIC_ID, 1), Set.of("rack2")
+        );
+
+        assertFalse(AssignorHelpers.useRackAwareAssignment(allMemberRacks, allPartitionRacks, racksPerPartition));
+    }
+
+    @Test
+    void testUseRackAwareAssignmentWithDisjointRacks() {
+        Set<String> allMemberRacks = Set.of("rack1", "rack2");
+        Set<String> allPartitionRacks = Set.of("rack3", "rack4");
+        Map<TopicIdPartition, Set<String>> racksPerPartition = Map.of(
+            new TopicIdPartition(TOPIC_ID, 0), Set.of("rack3"),
+            new TopicIdPartition(TOPIC_ID, 1), Set.of("rack4")
+        );
+
+        assertFalse(AssignorHelpers.useRackAwareAssignment(allMemberRacks, allPartitionRacks, racksPerPartition));
+    }
+
+    @Test
+    void testUseRackAwareAssignmentWithAllPartitionsHavingSameRacks() {
+        Set<String> allMemberRacks = Set.of("rack1", "rack2");
+        Set<String> allPartitionRacks = Set.of("rack1", "rack2");
+        Map<TopicIdPartition, Set<String>> racksPerPartition = Map.of(
+            new TopicIdPartition(TOPIC_ID, 0), Set.of("rack1", "rack2"),
+            new TopicIdPartition(TOPIC_ID, 1), Set.of("rack1", "rack2")
+        );
+
+        assertFalse(AssignorHelpers.useRackAwareAssignment(allMemberRacks, allPartitionRacks, racksPerPartition));
+    }
+
+    @Test
+    void testUseRackAwareAssignmentWithDifferentRacksPerPartition() {
+        Set<String> allMemberRacks = Set.of("rack1", "rack2");
+        Set<String> allPartitionRacks = Set.of("rack1", "rack2");
+        Map<TopicIdPartition, Set<String>> racksPerPartition = Map.of(
+            new TopicIdPartition(TOPIC_ID, 0), Set.of("rack1"),
+            new TopicIdPartition(TOPIC_ID, 1), Set.of("rack2")
+        );
+
+        assertTrue(AssignorHelpers.useRackAwareAssignment(allMemberRacks, allPartitionRacks, racksPerPartition));
+    }
+}