KAFKA-14702: Extend server side assignor to support rack aware replica placement

checkstyle/suppressions.xml

@@ -321,13 +321,13 @@
 
     <!-- group coordinator -->
     <suppress checks="CyclomaticComplexity"
+              files="(ConsumerGroupMember).java"/>
+    <suppress checks="ParameterNumber"
               files="(ConsumerGroupMember|GroupMetadataManager).java"/>
     <suppress checks="(NPathComplexity|MethodLength)"
               files="(GroupMetadataManager|ConsumerGroupTest|GroupMetadataManagerTest).java"/>
     <suppress checks="ClassFanOutComplexity"
               files="(GroupMetadataManager|GroupMetadataManagerTest).java"/>
-    <suppress checks="ParameterNumber"
-              files="(ConsumerGroupMember|GroupMetadataManager).java"/>
     <suppress checks="ClassDataAbstractionCouplingCheck"
               files="(RecordHelpersTest|GroupMetadataManagerTest).java"/>
     <suppress checks="JavaNCSS"

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java

@@ -777,7 +777,8 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
             subscriptionMetadata = group.computeSubscriptionMetadata(
                 member,
                 updatedMember,
-                metadataImage.topics()
+                metadataImage.topics(),
+                metadataImage.cluster()
             );
 
             if (!subscriptionMetadata.equals(group.subscriptionMetadata())) {
@@ -927,7 +928,8 @@ private List<Record> consumerGroupFenceMember(
         Map<String, TopicMetadata> subscriptionMetadata = group.computeSubscriptionMetadata(
             member,
             null,
-            metadataImage.topics()
+            metadataImage.topics(),
+            metadataImage.cluster()
         );
 
         if (!subscriptionMetadata.equals(group.subscriptionMetadata())) {

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java

@@ -131,13 +131,21 @@ public static Record newGroupSubscriptionMetadataRecord(
         Map<String, TopicMetadata> newSubscriptionMetadata
     ) {
         ConsumerGroupPartitionMetadataValue value = new ConsumerGroupPartitionMetadataValue();
-        newSubscriptionMetadata.forEach((topicName, topicMetadata) ->
+        newSubscriptionMetadata.forEach((topicName, topicMetadata) -> {
+            List<ConsumerGroupPartitionMetadataValue.PartitionMetadata> partitionMetadata = new ArrayList<>();
+            topicMetadata.partitionRacks().forEach((partition, racks) ->
+                partitionMetadata.add(new ConsumerGroupPartitionMetadataValue.PartitionMetadata()
+                    .setPartition(partition)
+                    .setRacks(new ArrayList<>(racks))
+                )
+            );
             value.topics().add(new ConsumerGroupPartitionMetadataValue.TopicMetadata()
                 .setTopicId(topicMetadata.id())
                 .setTopicName(topicMetadata.name())
                 .setNumPartitions(topicMetadata.numPartitions())
-            )
-        );
+                .setPartitionMetadata(partitionMetadata)
+            );
+        });
 
         return new Record(
             new ApiMessageAndVersion(

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

@@ -16,8 +16,6 @@
  */
 package org.apache.kafka.coordinator.group.assignor;
 
-import org.apache.kafka.common.Uuid;
-
 import java.util.Map;
 import java.util.Objects;
 
@@ -26,59 +24,39 @@
  */
 public class AssignmentSpec {
     /**
-     * The members keyed by member id.
+     * The members keyed by member Id.
      */
     private final Map<String, AssignmentMemberSpec> members;
 
-    /**
-     * The topics' metadata keyed by topic id.
-     */
-    private final Map<Uuid, AssignmentTopicMetadata> topics;
-
     public AssignmentSpec(
-        Map<String, AssignmentMemberSpec> members,
-        Map<Uuid, AssignmentTopicMetadata> topics
+        Map<String, AssignmentMemberSpec> members
     ) {
         Objects.requireNonNull(members);
-        Objects.requireNonNull(topics);
         this.members = members;
-        this.topics = topics;
     }
 
     /**
-     * @return Member metadata keyed by member Ids.
+     * @return Member metadata keyed by member Id.
      */
     public Map<String, AssignmentMemberSpec> members() {
         return members;
     }
 
-    /**
-     * @return Topic metadata keyed by topic Ids.
-     */
-    public Map<Uuid, AssignmentTopicMetadata> topics() {
-        return topics;
-    }
-
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         AssignmentSpec that = (AssignmentSpec) o;
-        if (!members.equals(that.members)) return false;
-        return topics.equals(that.topics);
+        return members.equals(that.members);
     }
 
     @Override
     public int hashCode() {
-        int result = members.hashCode();
-        result = 31 * result + topics.hashCode();
-        return result;
+        return members.hashCode();
     }
 
     @Override
     public String toString() {
-        return "AssignmentSpec(members=" + members +
-            ", topics=" + topics +
-            ')';
+        return "AssignmentSpec(members=" + members + ')';
     }
 }

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

@@ -36,8 +36,9 @@ public interface PartitionAssignor {
      * Perform the group assignment given the current members and
      * topic metadata.
      *
-     * @param assignmentSpec The assignment spec.
+     * @param assignmentSpec           The member assignment spec.
+     * @param subscribedTopicDescriber The topic and cluster metadata describer {@link SubscribedTopicDescriber}.
      * @return The new assignment for the group.
      */
-    GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAssignorException;
+    GroupAssignment assign(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) throws PartitionAssignorException;
 }

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

@@ -78,23 +78,19 @@ public MemberWithRemainingAssignments(String memberId, int remaining) {
     /**
      * @return Map of topic ids to a list of members subscribed to them.
      */
-    private Map<Uuid, List<String>> membersPerTopic(final AssignmentSpec assignmentSpec) {
+    private Map<Uuid, List<String>> membersPerTopic(final AssignmentSpec assignmentSpec, final SubscribedTopicDescriber subscribedTopicDescriber) {
         Map<Uuid, List<String>> membersPerTopic = new HashMap<>();
         Map<String, AssignmentMemberSpec> membersData = assignmentSpec.members();
 
         membersData.forEach((memberId, memberMetadata) -> {
             Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
             for (Uuid topicId : topics) {
-                // Only topics that are present in both the subscribed topics list and the topic metadata should be
-                // considered for assignment.
-                if (assignmentSpec.topics().containsKey(topicId)) {
-                    membersPerTopic
-                        .computeIfAbsent(topicId, k -> new ArrayList<>())
-                        .add(memberId);
-                } else {
-                    throw new PartitionAssignorException("Member " + memberId + " subscribed to topic " +
-                        topicId + " which doesn't exist in the topic metadata");
+                if (subscribedTopicDescriber.numPartitions(topicId) == -1) {
+                    throw new PartitionAssignorException("Member is subscribed to a non-existent topic");
                 }
+                membersPerTopic
+                    .computeIfAbsent(topicId, k -> new ArrayList<>())
+                    .add(memberId);
             }
         });
 
@@ -118,14 +114,14 @@ private Map<Uuid, List<String>> membersPerTopic(final AssignmentSpec assignmentS
      * </ol>
      */
     @Override
-    public GroupAssignment assign(final AssignmentSpec assignmentSpec) throws PartitionAssignorException {
+    public GroupAssignment assign(final AssignmentSpec assignmentSpec, final SubscribedTopicDescriber subscribedTopicDescriber) throws PartitionAssignorException {
         Map<String, MemberAssignment> newAssignment = new HashMap<>();
 
         // Step 1
-        Map<Uuid, List<String>> membersPerTopic = membersPerTopic(assignmentSpec);
+        Map<Uuid, List<String>> membersPerTopic = membersPerTopic(assignmentSpec, subscribedTopicDescriber);
 
         membersPerTopic.forEach((topicId, membersForTopic) -> {
-            int numPartitionsForTopic = assignmentSpec.topics().get(topicId).numPartitions();
+            int numPartitionsForTopic = subscribedTopicDescriber.numPartitions(topicId);
             int minRequiredQuota = numPartitionsForTopic / membersForTopic.size();
             // Each member can get only ONE extra partition per topic after receiving the minimum quota.
             int numMembersWithExtraPartition = numPartitionsForTopic % membersForTopic.size();

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

@@ -0,0 +1,51 @@
+/*
+ * 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.common.annotation.InterfaceStability;
+
+import java.util.Set;
+
+/**
+ * The subscribed topic describer is used by the {@link PartitionAssignor}
+ * to obtain topic and partition metadata of subscribed topics.
+ *
+ * The interface is kept in an internal module until KIP-848 is fully
+ * implemented and ready to be released.
+ */
+@InterfaceStability.Unstable
+public interface SubscribedTopicDescriber {
+    /**
+     * Number of partitions for the given topic Id.
+     *
+     * @param topicId   Uuid corresponding to the topic.
+     * @return The number of partitions corresponding to the given topicId.
+     *         If the topicId doesn't exist return -1;
+     */
+    int numPartitions(Uuid topicId);
+
+    /**
+     * Returns all the racks associated with the replicas for the given partition.
+     *
+     * @param topicId   Uuid corresponding to the partition's topic.
+     * @param partition Partition number within topic.
+     * @return The set of racks corresponding to the replicas of the topics partition.
+     *         If the topicId doesn't exist return an empty set;
+     */
+    Set<String> racksForPartition(Uuid topicId, int partition);
+}

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroup.java

@@ -19,6 +19,7 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.errors.UnknownMemberIdException;
 import org.apache.kafka.coordinator.group.Group;
+import org.apache.kafka.image.ClusterImage;
 import org.apache.kafka.image.TopicImage;
 import org.apache.kafka.image.TopicsImage;
 import org.apache.kafka.timeline.SnapshotRegistry;
@@ -28,6 +29,7 @@
 
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -434,22 +436,39 @@ public void setSubscriptionMetadata(
     public Map<String, TopicMetadata> computeSubscriptionMetadata(
         ConsumerGroupMember oldMember,
         ConsumerGroupMember newMember,
-        TopicsImage topicsImage
+        TopicsImage topicsImage,
+        ClusterImage clusterImage
     ) {
         // Copy and update the current subscriptions.
         Map<String, Integer> subscribedTopicNames = new HashMap<>(this.subscribedTopicNames);
         maybeUpdateSubscribedTopicNames(subscribedTopicNames, oldMember, newMember);
 
         // Create the topic metadata for each subscribed topic.
         Map<String, TopicMetadata> newSubscriptionMetadata = new HashMap<>(subscribedTopicNames.size());
+
         subscribedTopicNames.forEach((topicName, count) -> {
             TopicImage topicImage = topicsImage.getTopic(topicName);
             if (topicImage != null) {
+                Map<Integer, Set<String>> partitionRacks = new HashMap<>();
+
+                topicImage.partitions().forEach((partition, partitionRegistration) -> {
+                    Set<String> racks = new HashSet<>();
+                    for (int replica : partitionRegistration.replicas) {
+                        Optional<String> rackOptional = clusterImage.broker(replica).rack();
+                        // Only add rack if it is available for the broker/replica.
+                        rackOptional.ifPresent(racks::add);
+                    }
+                    // If no racks are available for any replica of this partition, store an empty map.
+                    if (!racks.isEmpty())
+                        partitionRacks.put(partition, racks);
+                });
+
                 newSubscriptionMetadata.put(topicName, new TopicMetadata(
                     topicImage.id(),
                     topicImage.name(),
-                    topicImage.partitions().size()
-                ));
+                    topicImage.partitions().size(),
+                    partitionRacks)
+                );
             }
         });