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

checkstyle/import-control.xml

@@ -48,6 +48,7 @@
   <allow pkg="org.apache.kafka.common.utils" />
   <allow pkg="org.apache.kafka.common.errors" exact-match="true" />
   <allow pkg="org.apache.kafka.common.memory" />
+  <allow pkg="org.apache.kafka.coordinator.group.common"/>
 
   <subpackage name="common">
     <allow class="org.apache.kafka.clients.consumer.ConsumerRecord" exact-match="true" />

checkstyle/suppressions.xml

@@ -321,6 +321,8 @@
 
     <!-- group coordinator -->
     <suppress checks="CyclomaticComplexity"
+              files="(ConsumerGroupMember).java"/>
+    <suppress checks="ParameterNumber"
               files="(ConsumerGroupMember|GroupMetadataManager).java"/>
     <suppress checks="MethodLength"
               files="(GroupMetadataManager|ConsumerGroupTest|GroupMetadataManagerTest).java"/>

clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java

@@ -131,8 +131,8 @@ private boolean allSubscriptionsEqual(Set<String> allTopics,
 
         Set<String> subscribedTopics = new HashSet<>();
 
-        // keep track of all previously owned partitions so we can invalidate them if invalid input is
-        // detected, eg two consumers somehow claiming the same partition in the same/current generation
+        // keep track of all previously owned partitions, so we can invalidate them if invalid input is
+        // detected, e.g. two consumers somehow claiming the same partition in the same/current generation
         Map<TopicPartition, String> allPreviousPartitionsToOwner = new HashMap<>();
 
         for (Map.Entry<String, Subscription> subscriptionEntry : subscriptions.entrySet()) {
@@ -361,9 +361,9 @@ private boolean hasCycles(Set<ConsumerPair> pairs) {
                 }
             }
 
-            // for now we want to make sure there is no partition movements of the same topic between a pair of consumers.
+            // for now, we want to make sure there is no partition movements of the same topic between a pair of consumers.
             // the odds of finding a cycle among more than two consumers seem to be very low (according to various randomized
-            // tests with the given sticky algorithm) that it should not worth the added complexity of handling those cases.
+            // tests with the given sticky algorithm) that it should not be worth the added complexity of handling those cases.
             for (List<String> cycle: cycles)
                 if (cycle.size() == 3) // indicates a cycle of length 2
                     return true;
@@ -806,7 +806,7 @@ private int assignNewPartition(TopicPartition unassignedPartition, String consum
         private void verifyUnfilledMembers() {
 
             if (!unfilledMembersWithUnderMinQuotaPartitions.isEmpty()) {
-                // we expected all the remaining unfilled members have minQuota partitions and we're already at the expected number
+                // we expected all the remaining unfilled members have minQuota partitions, and we're already at the expected number
                 // of members with more than the minQuota partitions. Otherwise, there must be error here.
                 if (currentNumMembersWithOverMinQuotaPartitions != expectedNumMembersWithOverMinQuotaPartitions) {
                     log.error("Current number of members with more than the minQuota partitions: {}, is less than the expected number " +
@@ -896,7 +896,7 @@ private List<TopicPartition> getUnassignedPartitions(List<TopicPartition> sorted
      *    while preserving rack-alignment. This step is used only for rack-aware assignment.
      * 4. Assigning the remaining unassigned partitions in a way that balances out the overall assignments of partitions to consumers.
      *    For rack-aware algorithm, these are partitions that could not be aligned on racks within the balancing constraints.
-     * 5. Further balancing out the resulting assignment by finding the partitions that can be reassigned
+     * 5. Further, balancing out the resulting assignment by finding the partitions that can be reassigned
      *    to another consumer towards an overall more balanced assignment. For rack-aware algorithm, attempt
      *    to retain rack alignment if possible.
      *
@@ -1021,7 +1021,7 @@ private List<TopicPartition> assignOwnedPartitions() {
                             revocationRequired = true;
                         } else {
                             // otherwise, remove the topic partition from those that need to be assigned only if
-                            // its current consumer is still subscribed to its topic (because it is already assigned
+                            // its current consumer is still subscribed to its topic (because it is already assigned,
                             // and we would want to preserve that assignment as much as possible)
                             assignedPartitions.add(partition);
                         }
@@ -1075,7 +1075,7 @@ private List<TopicPartition> getUnassignedPartitions(List<TopicPartition> sorted
          * update the prevAssignment with the partitions, consumer and generation in parameters
          *
          * @param partitions:       The partitions to be updated the prevAssignment
-         * @param consumer:         The consumer Id
+         * @param consumer:         The consumerId
          * @param prevAssignment:   The assignment contains the assignment with the 2nd largest generation
          * @param generation:       The generation of this assignment (partitions)
          */
@@ -1321,7 +1321,7 @@ private void balance(Map<TopicPartition, ConsumerGenerationPair> prevAssignment,
                     fixedAssignments.put(consumer, currentAssignment.remove(consumer));
                 }
 
-            // create a deep copy of the current assignment so we can revert to it if we do not get a more balanced assignment later
+            // create a deep copy of the current assignment, so we can revert to it if we do not get a more balanced assignment later
             Map<String, List<TopicPartition>> preBalanceAssignment = deepCopy(currentAssignment);
             Map<TopicPartition, String> preBalancePartitionConsumers = new HashMap<>(currentPartitionConsumer);
 
@@ -1332,7 +1332,7 @@ private void balance(Map<TopicPartition, ConsumerGenerationPair> prevAssignment,
 
             boolean reassignmentPerformed = performReassignments(sortedAllPartitions, prevAssignment);
 
-            // if we are not preserving existing assignments and we have made changes to the current assignment
+            // if we are not preserving existing assignments, and we have made changes to the current assignment
             // make sure we are getting a more balanced assignment; otherwise, revert to previous assignment
             if (!initializing && reassignmentPerformed && getBalanceScore(currentAssignment) >= getBalanceScore(preBalanceAssignment)) {
                 deepCopy(preBalanceAssignment, currentAssignment);
@@ -1358,7 +1358,7 @@ private boolean performReassignments(List<TopicPartition> reassignablePartitions
             // repeat reassignment until no partition can be moved to improve the balance
             do {
                 modified = false;
-                // reassign all reassignable partitions (starting from the partition with least potential consumers and if needed)
+                // reassign all reassignable partitions (starting from the partition with the least potential consumers and if needed)
                 // until the full list is processed or a balance is achieved
                 Iterator<TopicPartition> partitionIterator = reassignablePartitions.iterator();
                 while (partitionIterator.hasNext() && !isBalanced()) {

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

@@ -578,7 +578,8 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
             subscriptionMetadata = group.computeSubscriptionMetadata(
                 member,
                 updatedMember,
-                metadataImage.topics()
+                metadataImage.topics(),
+                metadataImage.cluster()
             );
 
             if (!subscriptionMetadata.equals(group.subscriptionMetadata())) {
@@ -609,7 +610,7 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
 
             try {
                 TargetAssignmentBuilder.TargetAssignmentResult assignmentResult =
-                    new TargetAssignmentBuilder(groupId, groupEpoch, assignors.get(preferredServerAssignor))
+                    new TargetAssignmentBuilder(groupId, groupEpoch,assignors.get(preferredServerAssignor))
                         .withMembers(group.members())
                         .withSubscriptionMetadata(subscriptionMetadata)
                         .withTargetAssignment(group.targetAssignment())
@@ -700,7 +701,8 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
         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/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;
 
@@ -30,19 +28,11 @@ public class AssignmentSpec {
      */
     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;
     }
 
     /**
@@ -52,33 +42,22 @@ 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;
+        if (!(o instanceof AssignmentSpec)) 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 Objects.hash(members);
     }
 
     @Override
     public String toString() {
         return "AssignmentSpec(members=" + members +
-            ", topics=" + topics +
             ')';
     }
 }

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

@@ -0,0 +1,59 @@
+/*
+ * 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 assignment 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 AssignmentTopicDescriber {
+
+    /**
+     * Returns a set of subscribed topicIds.
+     *
+     * @return Set of topicIds corresponding to the subscribed topics.
+     */
+    Set<Uuid> subscribedTopicIds();
+
+    /**
+     * Number of partitions for the given topicId.
+     *
+     * @param topicId   Uuid corresponding to the topic.
+     * @return The number of partitions corresponding to the given topicId.
+     *         If the topicId doesn't exist return 0;
+     */
+    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/assignor/AssignmentTopicMetadata.java

@@ -16,44 +16,83 @@
  */
 package org.apache.kafka.coordinator.group.assignor;
 
-/**
- * Metadata of a topic.
- */
-public class AssignmentTopicMetadata {
+import org.apache.kafka.common.Uuid;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+public class AssignmentTopicMetadata implements AssignmentTopicDescriber {
+
+    Map<Uuid, PartitionMetadata> TopicPartitionMetadata;
+
+    public AssignmentTopicMetadata(Map<Uuid, PartitionMetadata> TopicPartitionMetadata) {
+        this.TopicPartitionMetadata = TopicPartitionMetadata;
+    }
+
+    /**
+     * Returns a set of subscribed topicIds.
+     *
+     * @return Set of topicIds corresponding to the subscribed topics.
+     */
+    @Override
+    public Set<Uuid> subscribedTopicIds() {
+        return TopicPartitionMetadata.keySet();
+    }
 
     /**
-     * The number of partitions.
+     * Number of partitions for the given topicId.
+     *
+     * @param topicId   Uuid corresponding to the topic.
+     * @return The number of partitions corresponding to the given topicId.
+     *         If the topicId doesn't exist return 0;
      */
-    private final int numPartitions;
+    @Override
+    public int numPartitions(Uuid topicId) {
+        PartitionMetadata partitionMetadata = TopicPartitionMetadata.get(topicId);
+        if (partitionMetadata == null) {
+            return 0;
+        }
 
-    public AssignmentTopicMetadata(
-        int numPartitions
-    ) {
-        this.numPartitions = numPartitions;
+        return partitionMetadata.numPartitions();
     }
 
     /**
-     * @return The number of partitions present for the topic.
+     * 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;
      */
-    public int numPartitions() {
-        return numPartitions;
+    @Override
+    public Set<String> racksForPartition(Uuid topicId, int partition) {
+        PartitionMetadata partitionMetadata = TopicPartitionMetadata.get(topicId);
+        if (partitionMetadata == null) {
+            return Collections.emptySet();
+        }
+
+        return partitionMetadata.racks(partition);
     }
 
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
+        if (!(o instanceof AssignmentTopicMetadata)) return false;
         AssignmentTopicMetadata that = (AssignmentTopicMetadata) o;
-        return numPartitions == that.numPartitions;
+        return TopicPartitionMetadata.equals(that.TopicPartitionMetadata);
     }
 
     @Override
     public int hashCode() {
-        return numPartitions;
+        return Objects.hash(TopicPartitionMetadata);
     }
 
     @Override
     public String toString() {
-        return "AssignmentTopicMetadata(numPartitions=" + numPartitions + ')';
+        return "AssignmentTopicMetadata{" +
+            "TopicPartitionMetadata=" + TopicPartitionMetadata +
+            '}';
     }
 }

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 assignmentTopicDescriber The topic and cluster metadata describer {@link AssignmentTopicDescriber}.
+     * @param assignmentSpec The member assignment spec.
      * @return The new assignment for the group.
      */
-    GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAssignorException;
+    GroupAssignment assign(AssignmentTopicDescriber assignmentTopicDescriber, AssignmentSpec assignmentSpec) throws PartitionAssignorException;
 }