You are viewing a plain text version of this content. The canonical link for it is here.
Posted to jira@kafka.apache.org by "rreddy-22 (via GitHub)" <gi...@apache.org> on 2023/10/04 03:45:28 UTC
[PR] KAFKA-15484/General Rack Aware Assignor [kafka]
rreddy-22 opened a new pull request, #14481:
URL: https://github.com/apache/kafka/pull/14481
This is part two of the Uniform Assignor, used when the subscriptions of each member in a consumer group are different.
Files added are:
GeneralUniformAssignmentBuilder: Contains all the assignment logic for the assignor.
GeneralUniformAssignmentBuilderTest: Contains tests for most scenarios that could cause a rebalance.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401000055
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) || !membersPerTopic.containsKey(reassignablePartition.topicId()))
+ continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember : sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ if (newOwner == null) {
+ throw new PartitionAssignorException("No suitable new owner was found for the partition" + partition);
+ }
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromTargetAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToTargetAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberPair(" +
+ "srcMemberId='" + srcMemberId + '\'' +
+ ", dstMemberId='" + dstMemberId + '\'' +
+ ')';
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among members.
+ * During a partition rebalance, it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
Review Comment:
yes that's precisely it, just to make sure we're not moving partitions back and forth mainly.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1400998694
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformAssignor.java:
##########
@@ -96,7 +96,7 @@ public GroupAssignment assign(
private boolean allSubscriptionsEqual(Map<String, AssignmentMemberSpec> members) {
Set<Uuid> firstSubscriptionSet = new HashSet<>(members.values().iterator().next().subscribedTopicIds());
for (AssignmentMemberSpec memberSpec : members.values()) {
- if (!firstSubscriptionSet.containsAll(memberSpec.subscribedTopicIds())) {
+ if (!firstSubscriptionSet.equals(new HashSet<>(memberSpec.subscribedTopicIds()))) {
Review Comment:
yes cause the old way means every element in memberSpec.subscribedTopicIds() must be in firstSubscriptionSet, but firstSubscriptionSet can have additional elements that are not in memberSpec.subscribedTopicIds(). We need a strict equality check
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396618526
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
Review Comment:
ohh thanks for bringing it up, I wasn't aware of this, I was following the existing code from the client side assignor. Let me see what I can change it to
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396639777
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
+ return false;
+ }
+
+ addPartitionToMembersAssignment(topicIdPartition, memberId);
+ return true;
+ }
+
+ /**
+ * Assigns a partition to a member (if eligible), updates the current assignment size,
Review Comment:
oh yes you're right, I forgot to change the javadoc after changing the method, initially I implemented it differently. thanks for the catch!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "jeffkbkim (via GitHub)" <gi...@apache.org>.
jeffkbkim commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396385183
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
Review Comment:
should we pass in all of the global variables used in `initializeAssignmentSizes`?
also, i don't see a reason we have a helper method
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
Review Comment:
intuitively, it makes sense to me we only add the partition to `assignedStickyPartitions` if the partition was actually assigned to the member. why don't we check the return value of `maybeAssignPartitionToMember` here?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
Review Comment:
what does "on the client" mean?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
Review Comment:
nit: "During a partition rebalance, it keeps track..."
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
Review Comment:
`if (assignmentManager.isMemberAtCapacity(memberId))` or something along those lines would simplify other parts of the code as well.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
Review Comment:
nit: members
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
Review Comment:
this method sorts twice across all partitions. might need to look into this
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
+ return false;
+ }
+
+ addPartitionToMembersAssignment(topicIdPartition, memberId);
+ return true;
+ }
+
+ /**
+ * Assigns a partition to a member (if eligible), updates the current assignment size,
+ * and updates relevant data structures.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId Member that the partition needs to be added to.
+ */
+ private void addPartitionToMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
+ addPartitionToAssignment(
+ targetAssignment,
+ memberId,
+ topicIdPartition.topicId(),
+ topicIdPartition.partitionId()
+ );
+
+ partitionOwnerInTargetAssignment.put(topicIdPartition, memberId);
+ // Remove the member's assignment data from the queue to update it.
+ sortedMembersByAssignmentSize.remove(memberId);
+ assignmentManager.incrementTargetAssignmentSize(memberId);
+
+ // Update current assignment size and re-add to queue if needed.
+ if (assignmentManager.targetAssignmentSize(memberId) < assignmentManager.maxAssignmentSize(memberId)) {
+ sortedMembersByAssignmentSize.add(memberId);
+ }
+
+ unassignedPartitions.remove(topicIdPartition);
+ }
+
+ /**
+ * Revokes the partition from a member, updates the current target assignment size,
+ * and other relevant data structures.
+ *
+ * @param topicIdPartition The partition to be revoked.
+ * @param memberId Member that the partition needs to be revoked from.
+ */
+ private void removePartitionFromMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
+ Map<Uuid, Set<Integer>> targetPartitionsMap = targetAssignment.get(memberId).targetPartitions();
+ Set<Integer> partitionsSet = targetPartitionsMap.get(topicIdPartition.topicId());
+ // Remove the partition from the assignment, if there are no more partitions from a particular topic,
+ // remove the topic from the assignment as well.
+ if (partitionsSet != null) {
+ partitionsSet.remove(topicIdPartition.partitionId());
+ if (partitionsSet.isEmpty()) {
+ targetPartitionsMap.remove(topicIdPartition.topicId());
+ }
+ }
+
+ partitionOwnerInTargetAssignment.remove(topicIdPartition, memberId);
+ // Remove the member's assignment data from the queue to update it.
+ sortedMembersByAssignmentSize.remove(memberId);
+ assignmentManager.decrementTargetAssignmentSize(memberId);
+
+ // Update current assignment size and re-add to queue if needed.
+ if (assignmentManager.targetAssignmentSize(memberId) < assignmentManager.maxAssignmentSize(memberId)) {
+ sortedMembersByAssignmentSize.add(memberId);
+ }
+ }
+
+ /**
+ * Sorts members in ascending order based on their current target assignment size.
+ * Members that have reached their max assignment size are removed.
+ *
+ * @param memberIds All the member Ids that need to be sorted.
Review Comment:
nit: "All of the". Actually, "The member Ids..." is simpler
i notice we use "all the" in a bunch of places, let's fix that
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
+ return false;
+ }
+
+ addPartitionToMembersAssignment(topicIdPartition, memberId);
+ return true;
+ }
+
+ /**
+ * Assigns a partition to a member (if eligible), updates the current assignment size,
+ * and updates relevant data structures.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId Member that the partition needs to be added to.
+ */
+ private void addPartitionToMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
+ addPartitionToAssignment(
Review Comment:
addPartitionToTargetAssignment to disambiguate. seems to imply we could be assigning to other assignments
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
+ return false;
+ }
+
+ addPartitionToMembersAssignment(topicIdPartition, memberId);
+ return true;
+ }
+
+ /**
+ * Assigns a partition to a member (if eligible), updates the current assignment size,
+ * and updates relevant data structures.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId Member that the partition needs to be added to.
+ */
+ private void addPartitionToMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
+ addPartitionToAssignment(
+ targetAssignment,
+ memberId,
+ topicIdPartition.topicId(),
+ topicIdPartition.partitionId()
+ );
+
+ partitionOwnerInTargetAssignment.put(topicIdPartition, memberId);
+ // Remove the member's assignment data from the queue to update it.
+ sortedMembersByAssignmentSize.remove(memberId);
+ assignmentManager.incrementTargetAssignmentSize(memberId);
+
+ // Update current assignment size and re-add to queue if needed.
+ if (assignmentManager.targetAssignmentSize(memberId) < assignmentManager.maxAssignmentSize(memberId)) {
+ sortedMembersByAssignmentSize.add(memberId);
+ }
Review Comment:
can we do something like
```
assignmentManager.incrementTargetAssignmentSize(memberId);
if (assignmentManager.isMemberAtCapacity(memberId)) {
sortedMembersByAssignmentSize.remove(memberId)
}
```
instead? also, this is technically not a queue, it's a set
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
+ return false;
+ }
+
+ addPartitionToMembersAssignment(topicIdPartition, memberId);
+ return true;
+ }
+
+ /**
+ * Assigns a partition to a member (if eligible), updates the current assignment size,
+ * and updates relevant data structures.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId Member that the partition needs to be added to.
+ */
+ private void addPartitionToMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
Review Comment:
nit: `addPartitionToTargetAssignment`
i think we should conform to using target assignment as the final assignment instead of using member(s) assignments
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
Review Comment:
it would be good if we didn't have to rebuild the topic partition set
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
+ return false;
+ }
+
+ addPartitionToMembersAssignment(topicIdPartition, memberId);
+ return true;
+ }
+
+ /**
+ * Assigns a partition to a member (if eligible), updates the current assignment size,
Review Comment:
it seems we always assign. which part makes it `(if eligible)`?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
Review Comment:
nit: can we add the curly braces? we do so in L285
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
Review Comment:
I would put a link to the actual method. I don't see any method called `assignPartitions`
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396667524
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
Review Comment:
this is my thought process do correct me if I'm wrong, I used public and protected so that the outer class can access the methods. B/w public and protected I used public for methods that could potentially be reused by other classes that need this assignment manager functionality but I realize this access is limited by the protected nature of the class itself.
I think for now I can make the class private and all the methods package private so that they can be accessed by the outer class? The access modifiers always confuse me, I'm open to suggestions in this matter! Thanks!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "dajac (via GitHub)" <gi...@apache.org>.
dajac commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1400581697
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AbstractUniformAssignmentBuilder.java:
##########
@@ -223,22 +223,13 @@ protected boolean racksMismatch(String memberId, TopicIdPartition tp) {
* @return A sorted list of partitions with potential members in the same rack.
*/
protected List<TopicIdPartition> sortPartitionsByRackMembers(Collection<TopicIdPartition> topicIdPartitions) {
-
- List<TopicIdPartition> sortedPartitionsList = topicIdPartitions.stream()
- .sorted(Comparator.comparing(TopicIdPartition::topicId).thenComparing(TopicIdPartition::partitionId))
+ return topicIdPartitions.stream()
+ .filter(tp -> membersWithSameRackAsPartition.containsKey(tp) && !membersWithSameRackAsPartition.get(tp).isEmpty())
+ .sorted(Comparator.comparing(
+ (TopicIdPartition tp) -> membersWithSameRackAsPartition.getOrDefault(tp, Collections.emptyList()).size())
Review Comment:
nit: I wonder if we could omit the `TopicIdPartition` type here.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
Review Comment:
nit: I suppose that we don't really need `topics` here.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
Review Comment:
nit: ditto about the type here.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
Review Comment:
I am a bit worried about the rebalance part, especially with regarding to how it will scale with the number of member and partitions in the group. Should we file a jira to look into this a bit more closely at some point? It seems that there are opportunities to make it perform better. For instance, we could try to make it incremental.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) || !membersPerTopic.containsKey(reassignablePartition.topicId()))
+ continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember : sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ if (newOwner == null) {
+ throw new PartitionAssignorException("No suitable new owner was found for the partition" + partition);
+ }
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromTargetAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToTargetAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberPair(" +
+ "srcMemberId='" + srcMemberId + '\'' +
+ ", dstMemberId='" + dstMemberId + '\'' +
+ ')';
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among members.
+ * During a partition rebalance, it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(
+ TopicIdPartition partition,
+ String oldOwner,
+ String newOwner
+ ) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ String expectedOldOwner = partitionMovementsByPartition.get(partition).dstMemberId;
+ if (!oldOwner.equals(expectedOldOwner)) {
+ throw new PartitionAssignorException("Old owner does not match expected value for partition: " + partition);
+ }
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
Review Comment:
nit: Do we actually need the member id here? It does not seem to be used anywhere.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
Review Comment:
+1 if possible.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformAssignor.java:
##########
@@ -96,7 +96,7 @@ public GroupAssignment assign(
private boolean allSubscriptionsEqual(Map<String, AssignmentMemberSpec> members) {
Set<Uuid> firstSubscriptionSet = new HashSet<>(members.values().iterator().next().subscribedTopicIds());
for (AssignmentMemberSpec memberSpec : members.values()) {
- if (!firstSubscriptionSet.containsAll(memberSpec.subscribedTopicIds())) {
+ if (!firstSubscriptionSet.equals(new HashSet<>(memberSpec.subscribedTopicIds()))) {
Review Comment:
Is it really better?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
Review Comment:
It is still there...
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
Review Comment:
i agree with Jeff. i also find the encapsulation a bit weird here.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
Review Comment:
Thanks. It would be great if we could add some javadoc to it. I also wonder if we should call it `computePartitionToBeMoved` or something like that.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401400649
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
Review Comment:
hey I understand your concern but with every perform reassignments iteration we have to go through all the reassignable partitions again, it's not the case that once a partition is reassigned then we don't check it again sadly. The sorting is done to reduce the total number of iterations we're trying to do. We only sort once per iteration. Let me know if there are any other concerns or things I've missed or if you think this is reasonable
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401429251
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
Review Comment:
https://github.com/apache/kafka/pull/14481#discussion_r1401419710
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1400990586
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
Review Comment:
Could you please elaborate more on making it incremental?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396617558
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
Review Comment:
I see where you are coming from, I was actually hoping to avoid the method call and local vars initialization altogether if it's balanced, not sure how much it really helps in terms of optimization, what do you think? I also thought it would make logical sense immediately for anyone reading
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "dajac (via GitHub)" <gi...@apache.org>.
dajac commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1394378983
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
Review Comment:
It seems that we iterate over the members a few times to initialize a few things (L128, L141, L173). I wonder if we could consolidate all of them in one iteration. For instance, we could do it at the end of the constructor. What do you think?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
Review Comment:
nit: Should we just do this in the constructor as well?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
Review Comment:
nit: There is an extra space after the `if`.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
Review Comment:
nit: I would put the last closing parenthesis on a new line.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
Review Comment:
nit: Add `@Override`. We usually also put toString last in the class.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
Review Comment:
ditto.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
Review Comment:
ditto about `assert`.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
Review Comment:
nit: subscribedTopicIds?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
Review Comment:
nit: In other such loops, you had a space before the `:`. Actually, you use both in the file. I would just pick one and stay consistent.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
Review Comment:
nit: We could remove this empty line.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
Review Comment:
Should we put those two right after `int memberPartitionCount = ...`?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
Review Comment:
nit: This reads a bit weird. I would just keep the condition on one line.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
Review Comment:
Do we really need to create a ArrayList from the Set?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
Review Comment:
Could we actually return in this case to break both loops?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
Review Comment:
When `performReassignments` is called for the first time, we check this. However, we also just checked it as part at L400 as well. Should we just remove it there?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
Review Comment:
ditto.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
Review Comment:
Do we need this? It would be better to throw an IllegalStateException if we do. We usually don't use assert in the code base.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
Review Comment:
I have a hard time to really understand what this method does. Could you help me out?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
Review Comment:
nit: We usually use `()` instead of `{}`.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
Review Comment:
Why is the class `protected`?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
Review Comment:
I have noticed that we use have both `public` and `protected` methods in this class and I am not sure to understand why. Could you help me to understand this?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
+ return false;
+ }
+
+ addPartitionToMembersAssignment(topicIdPartition, memberId);
+ return true;
+ }
+
+ /**
+ * Assigns a partition to a member (if eligible), updates the current assignment size,
+ * and updates relevant data structures.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId Member that the partition needs to be added to.
+ */
+ private void addPartitionToMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
+ addPartitionToAssignment(
+ targetAssignment,
+ memberId,
+ topicIdPartition.topicId(),
+ topicIdPartition.partitionId()
+ );
+
+ partitionOwnerInTargetAssignment.put(topicIdPartition, memberId);
+ // Remove the member's assignment data from the queue to update it.
+ sortedMembersByAssignmentSize.remove(memberId);
+ assignmentManager.incrementTargetAssignmentSize(memberId);
+
+ // Update current assignment size and re-add to queue if needed.
+ if (assignmentManager.targetAssignmentSize(memberId) < assignmentManager.maxAssignmentSize(memberId)) {
+ sortedMembersByAssignmentSize.add(memberId);
+ }
+
+ unassignedPartitions.remove(topicIdPartition);
+ }
+
+ /**
+ * Revokes the partition from a member, updates the current target assignment size,
+ * and other relevant data structures.
+ *
+ * @param topicIdPartition The partition to be revoked.
+ * @param memberId Member that the partition needs to be revoked from.
+ */
+ private void removePartitionFromMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
+ Map<Uuid, Set<Integer>> targetPartitionsMap = targetAssignment.get(memberId).targetPartitions();
+ Set<Integer> partitionsSet = targetPartitionsMap.get(topicIdPartition.topicId());
+ // Remove the partition from the assignment, if there are no more partitions from a particular topic,
+ // remove the topic from the assignment as well.
+ if (partitionsSet != null) {
+ partitionsSet.remove(topicIdPartition.partitionId());
+ if (partitionsSet.isEmpty()) {
+ targetPartitionsMap.remove(topicIdPartition.topicId());
+ }
+ }
+
+ partitionOwnerInTargetAssignment.remove(topicIdPartition, memberId);
+ // Remove the member's assignment data from the queue to update it.
+ sortedMembersByAssignmentSize.remove(memberId);
+ assignmentManager.decrementTargetAssignmentSize(memberId);
+
+ // Update current assignment size and re-add to queue if needed.
+ if (assignmentManager.targetAssignmentSize(memberId) < assignmentManager.maxAssignmentSize(memberId)) {
+ sortedMembersByAssignmentSize.add(memberId);
+ }
+ }
+
+ /**
+ * Sorts members in ascending order based on their current target assignment size.
+ * Members that have reached their max assignment size are removed.
+ *
+ * @param memberIds All the member Ids that need to be sorted.
+ * @return A set that maintains the order of members by assignment size.
+ */
+ protected TreeSet<String> getSortedMembersByAssignmentSize(Collection<String> memberIds) {
Review Comment:
nit: `sortMembersByAssignmentSize`?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
Review Comment:
nit: Could we make it static?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396672967
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
Review Comment:
you're right, let me update the java docs, thanks
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396614486
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
Review Comment:
haha yes I was also concerned about the same thing and played around with a couple types of formatting, the line is too long is that okay if I keep it one line?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396613887
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
Review Comment:
yep that's true, we could do that, since it would go back to the start of the do while loop and break anyways. Thanks!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "jeffkbkim (via GitHub)" <gi...@apache.org>.
jeffkbkim commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1400069229
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
Review Comment:
nit: with the "smallest"
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
Review Comment:
nit: assigned to members
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
Review Comment:
This sentence is hard to follow, can we simplify it?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) || !membersPerTopic.containsKey(reassignablePartition.topicId()))
+ continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
Review Comment:
to confirm, we're saying we have the most balanced result once we have nothing to modify right? even it's still unbalanced
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
Review Comment:
nit: In `performReassignments()` we iterate through partitions/members that we can reassign to others. Whereas in `isBalanced()` we iterate through members that can receive partitions. I feel it would be simpler if they follow the same pattern. It's not super important, but i got confused while reading the code
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
Review Comment:
nit: two members
should we throw an exception here? we don't expect any partition to have less than 2 members since we already removed all of them before calling performReassignments() right?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
Review Comment:
this looks to have heavy performance implications. We start with a total N number of partitions and if we reassign a single rack matching partition on every iteration, we're doing N * N log N operations.
Instead, can we just remove the reassigned partition from `reassignablePartitions`? Then we'd still have the sorted reassignable partitions without having to re-sort this on every iteration.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) || !membersPerTopic.containsKey(reassignablePartition.topicId()))
+ continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember : sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ if (newOwner == null) {
+ throw new PartitionAssignorException("No suitable new owner was found for the partition" + partition);
+ }
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromTargetAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToTargetAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
Review Comment:
nit: how's MemberPath? to hint that there's a start and end. for me, the word pair does not bring much value to understanding what this class represents
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
Review Comment:
this is a bit confusing; we have `partitionOwnerInTargetAssignment` and `currentPartitionOwners`.
"The topic partition must have a target owner"
and `String targetOwner = ...`
is more readable to me. wdyt?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
Review Comment:
was this addressed?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) || !membersPerTopic.containsKey(reassignablePartition.topicId()))
+ continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember : sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ if (newOwner == null) {
+ throw new PartitionAssignorException("No suitable new owner was found for the partition" + partition);
+ }
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromTargetAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToTargetAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberPair(" +
+ "srcMemberId='" + srcMemberId + '\'' +
+ ", dstMemberId='" + dstMemberId + '\'' +
+ ')';
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among members.
+ * During a partition rebalance, it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
Review Comment:
david had a similar comment here: https://github.com/apache/kafka/pull/14481#discussion_r1394456067
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) || !membersPerTopic.containsKey(reassignablePartition.topicId()))
+ continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember : sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ if (newOwner == null) {
+ throw new PartitionAssignorException("No suitable new owner was found for the partition" + partition);
+ }
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromTargetAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToTargetAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberPair(" +
+ "srcMemberId='" + srcMemberId + '\'' +
+ ", dstMemberId='" + dstMemberId + '\'' +
+ ')';
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among members.
+ * During a partition rebalance, it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
Review Comment:
i'm still not sure why it's necessary to track the previous movements.
can you help me understand why we need this? Is this to prevent cyclic reassignments?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1397683148
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
Review Comment:
https://github.com/apache/kafka/pull/14481#discussion_r1394487935
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401404039
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) || !membersPerTopic.containsKey(reassignablePartition.topicId()))
+ continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember : sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ if (newOwner == null) {
+ throw new PartitionAssignorException("No suitable new owner was found for the partition" + partition);
+ }
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromTargetAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToTargetAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
Review Comment:
I understand your concern but since MemberPair is always used with src and dest I feel like it should be easy enough to comprehend esp with the java docs but if it seems too confusing I can change it. What do you think?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401413975
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
Review Comment:
makes sense, but it might also be better if we check balance differently from how we reassign the partitions, I can add it as a follow up if necessary
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401001513
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) || !membersPerTopic.containsKey(reassignablePartition.topicId()))
+ continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
Review Comment:
yep that's correct
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1400975524
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
Review Comment:
I had addressed removing it from the check before calling perform reassignments again, added to the condition now, thanks for pointing it out
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396638565
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
Review Comment:
yep sure, this method is determining which partition should actually be moved to a new owner, considering the history of partition movements and ownership changes. If a partition has been moved before or if there's a recorded reverse movement, it takes these into account to decide the next move. I can explain better offline if required.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401397868
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
Review Comment:
yes that's right, changed it to an exception
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401429553
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
Review Comment:
done thanks!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "dajac (via GitHub)" <gi...@apache.org>.
dajac merged PR #14481:
URL: https://github.com/apache/kafka/pull/14481
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396592116
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
Review Comment:
makes sense you're right, I initialized them all in one go now! Thankss for pointing it out
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396667524
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
Review Comment:
this is my thought process do correct me if I'm wrong, I used public and protected so that the outer class can access the methods. B/w public and protected I used public for methods that could potentially be reused by other classes that need this assignment manager functionality but I realize this access is limited by the protected nature of the class itself.
I think for now I can make the class private and all the methods public or package private so that they can be accessed by the outer class? The access modifiers always confuse me, I'm open to suggestions in this matter! Thanks!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396671729
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
+ return false;
+ }
+
+ addPartitionToMembersAssignment(topicIdPartition, memberId);
+ return true;
+ }
+
+ /**
+ * Assigns a partition to a member (if eligible), updates the current assignment size,
+ * and updates relevant data structures.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId Member that the partition needs to be added to.
+ */
+ private void addPartitionToMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
+ addPartitionToAssignment(
+ targetAssignment,
+ memberId,
+ topicIdPartition.topicId(),
+ topicIdPartition.partitionId()
+ );
+
+ partitionOwnerInTargetAssignment.put(topicIdPartition, memberId);
+ // Remove the member's assignment data from the queue to update it.
+ sortedMembersByAssignmentSize.remove(memberId);
+ assignmentManager.incrementTargetAssignmentSize(memberId);
+
+ // Update current assignment size and re-add to queue if needed.
+ if (assignmentManager.targetAssignmentSize(memberId) < assignmentManager.maxAssignmentSize(memberId)) {
+ sortedMembersByAssignmentSize.add(memberId);
+ }
Review Comment:
oh yes thanks for the catch, initially it was a queue and I forgot to change the comment!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396679975
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
Review Comment:
oh yes you're right, I've updated the method, thanks
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396675130
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
Review Comment:
yeah that's a valid concern, let me write a test for this and think about it, but we can probably add the check to be safe anyways. I think the conditions that we check for in maybeAssignPartitionsToMember are unlikely to be satisfied for this particular case that's why it wasn't checked.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401401639
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
Review Comment:
makes sense, I changed it to currentTargetOwner
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "jeffkbkim (via GitHub)" <gi...@apache.org>.
jeffkbkim commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1400044251
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
Review Comment:
we bring in dependencies - this constructor can only be built once the members and subscribedTopicDescriber global variables are initialized. This will break if that order changes (which we probably won't change). Since the constructor has no arguments it makes it seem as if it does not depend on anything. fwiw i don't think we have to change it, but just something to keep in mind
constructors are to initialize fields so i don't think we need a helper method
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1399988448
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
Review Comment:
discussed on dm
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401419710
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AbstractUniformAssignmentBuilder.java:
##########
@@ -223,22 +223,13 @@ protected boolean racksMismatch(String memberId, TopicIdPartition tp) {
* @return A sorted list of partitions with potential members in the same rack.
*/
protected List<TopicIdPartition> sortPartitionsByRackMembers(Collection<TopicIdPartition> topicIdPartitions) {
-
- List<TopicIdPartition> sortedPartitionsList = topicIdPartitions.stream()
- .sorted(Comparator.comparing(TopicIdPartition::topicId).thenComparing(TopicIdPartition::partitionId))
+ return topicIdPartitions.stream()
+ .filter(tp -> membersWithSameRackAsPartition.containsKey(tp) && !membersWithSameRackAsPartition.get(tp).isEmpty())
+ .sorted(Comparator.comparing(
+ (TopicIdPartition tp) -> membersWithSameRackAsPartition.getOrDefault(tp, Collections.emptyList()).size())
Review Comment:
If I omit the type then I get an error cause the topicId method can't be inferred correctly by the compiler, I need to explicitly mention the type. Do you have a workaround for this?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401404458
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) || !membersPerTopic.containsKey(reassignablePartition.topicId()))
+ continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember : sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ if (newOwner == null) {
+ throw new PartitionAssignorException("No suitable new owner was found for the partition" + partition);
+ }
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromTargetAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToTargetAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberPair(" +
+ "srcMemberId='" + srcMemberId + '\'' +
+ ", dstMemberId='" + dstMemberId + '\'' +
+ ')';
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among members.
+ * During a partition rebalance, it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
Review Comment:
ah F I missed this place, I'll change it, thanks for pointing it out
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396596218
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
Review Comment:
thanks, makes sense, I'll put a space before and after the colon in all the for loops
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396624092
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
Review Comment:
I threw a PartitionAssignorException for now, I'm actually not fully convinced anymore about throwing an exception even though it made sense to me back then, I'll double check this and get back to you if that's okay. If you have any thoughts on this I would appreciate that as well! thanks!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396670969
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
+ return false;
+ }
+
+ addPartitionToMembersAssignment(topicIdPartition, memberId);
+ return true;
+ }
+
+ /**
+ * Assigns a partition to a member (if eligible), updates the current assignment size,
+ * and updates relevant data structures.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId Member that the partition needs to be added to.
+ */
+ private void addPartitionToMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
Review Comment:
I agree, that would be consistent, thanks for pointing it out!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396669760
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
Review Comment:
thanks for pointing it out, I meant to be explicit about how the assignment is the one being used by the client currently, I can change it to member, what do you think?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396624783
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
Review Comment:
yes that's an oversight on my part thanks!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396670619
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
+ TopicIdPartition topicIdPartition,
+ String memberId
+ ) {
+ // If member is not subscribed to the partition's topic, return false without assigning.
+ if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
+ return false;
+ }
+
+ // If the member's current assignment is already at max, return false without assigning.
+ if (assignmentManager.targetAssignmentSize(memberId) >= assignmentManager.maxAssignmentSize(memberId)) {
+ return false;
+ }
+
+ addPartitionToMembersAssignment(topicIdPartition, memberId);
+ return true;
+ }
+
+ /**
+ * Assigns a partition to a member (if eligible), updates the current assignment size,
+ * and updates relevant data structures.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId Member that the partition needs to be added to.
+ */
+ private void addPartitionToMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
+ addPartitionToAssignment(
+ targetAssignment,
+ memberId,
+ topicIdPartition.topicId(),
+ topicIdPartition.partitionId()
+ );
+
+ partitionOwnerInTargetAssignment.put(topicIdPartition, memberId);
+ // Remove the member's assignment data from the queue to update it.
+ sortedMembersByAssignmentSize.remove(memberId);
+ assignmentManager.incrementTargetAssignmentSize(memberId);
+
+ // Update current assignment size and re-add to queue if needed.
+ if (assignmentManager.targetAssignmentSize(memberId) < assignmentManager.maxAssignmentSize(memberId)) {
+ sortedMembersByAssignmentSize.add(memberId);
+ }
+
+ unassignedPartitions.remove(topicIdPartition);
+ }
+
+ /**
+ * Revokes the partition from a member, updates the current target assignment size,
+ * and other relevant data structures.
+ *
+ * @param topicIdPartition The partition to be revoked.
+ * @param memberId Member that the partition needs to be revoked from.
+ */
+ private void removePartitionFromMembersAssignment(TopicIdPartition topicIdPartition, String memberId) {
+ Map<Uuid, Set<Integer>> targetPartitionsMap = targetAssignment.get(memberId).targetPartitions();
+ Set<Integer> partitionsSet = targetPartitionsMap.get(topicIdPartition.topicId());
+ // Remove the partition from the assignment, if there are no more partitions from a particular topic,
+ // remove the topic from the assignment as well.
+ if (partitionsSet != null) {
+ partitionsSet.remove(topicIdPartition.partitionId());
+ if (partitionsSet.isEmpty()) {
+ targetPartitionsMap.remove(topicIdPartition.topicId());
+ }
+ }
+
+ partitionOwnerInTargetAssignment.remove(topicIdPartition, memberId);
+ // Remove the member's assignment data from the queue to update it.
+ sortedMembersByAssignmentSize.remove(memberId);
+ assignmentManager.decrementTargetAssignmentSize(memberId);
+
+ // Update current assignment size and re-add to queue if needed.
+ if (assignmentManager.targetAssignmentSize(memberId) < assignmentManager.maxAssignmentSize(memberId)) {
+ sortedMembersByAssignmentSize.add(memberId);
+ }
+ }
+
+ /**
+ * Sorts members in ascending order based on their current target assignment size.
+ * Members that have reached their max assignment size are removed.
+ *
+ * @param memberIds All the member Ids that need to be sorted.
Review Comment:
I hear you, I will change it, thanks!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396669104
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
Review Comment:
could you help me understand why we should pass the global variables to initializeAssignmentSIzes?
I see your point about the helper method, my thought process was that I wanted it to be concise and easy to read.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "dajac (via GitHub)" <gi...@apache.org>.
dajac commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1397596576
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The current assignment size for the given member.
+ */
+ public int targetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.currentAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * @return The maximum assignment size for the given member.
+ */
+ public int maxAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return 0;
+ }
+ return memberData.maxAssignmentSize;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Increment the current target assignment size for the member.
+ */
+ public void incrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+ memberData.currentAssignmentSize++;
+ }
+
+ /**
+ * @param memberId The member Id.
+ * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
+ */
+ public void decrementTargetAssignmentSize(String memberId) {
+ MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
+ if (memberData == null) {
+ LOG.warn("Member Id {} not found", memberId);
+ return;
+ }
+
+ if (memberData.currentAssignmentSize > 0) {
+ memberData.currentAssignmentSize --;
+ }
+ }
+
+ /**
+ * Computes and sets the maximum assignment size for each member.
+ */
+ private void initializeAssignmentSizes() {
+ members.forEach((memberId, member) -> {
+ int maxSize = member.subscribedTopicIds().stream()
+ .mapToInt(subscribedTopicDescriber::numPartitions)
+ .sum();
+
+ MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
+ .computeIfAbsent(memberId, MemberAssignmentData::new);
+ memberAssignmentData.maxAssignmentSize = maxSize;
+ memberAssignmentData.currentAssignmentSize = 0;
+ });
+ }
+
+ /**
+ * Assigns partition to member if eligible.
+ *
+ * @param topicIdPartition The partition to be assigned.
+ * @param memberId The Id of the member.
+ * @return true if the partition was assigned, false otherwise.
+ */
+ protected boolean maybeAssignPartitionToMember(
Review Comment:
Yeah, keeping it internal for now makes sense to me.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401429874
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
Review Comment:
I actually just removed the check, it' not necessary
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
Review Comment:
I had addressed removing it from the check before calling perform reassignments again, added to the condition now, thanks for pointing it out
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "jeffkbkim (via GitHub)" <gi...@apache.org>.
jeffkbkim commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1400044251
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
Review Comment:
we bring in dependencies - this constructor can only be built once the members and subscribedTopicDescriber global variables are initialized. This will break if that order changes (which we probably won't change). Since the constructor has no arguments it makes it seem as if it does not depend on anything.
constructors are to initialize fields so i don't think we need a helper method
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401392351
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
+ "memberId='" + memberId + '\'' +
+ ", currentAssignmentSize=" + currentAssignmentSize +
+ ", maxAssignmentSize=" + maxAssignmentSize +
+ '}';
+ }
+ }
+
+ /**
+ * Initializes an AssignmentManager, setting up the necessary data structures.
+ */
+ public AssignmentManager() {
+ initializeAssignmentSizes();
+ }
Review Comment:
understood thanks for the explanation!
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401429114
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ return;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) || !membersPerTopic.containsKey(reassignablePartition.topicId()))
+ continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember : sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ if (newOwner == null) {
+ throw new PartitionAssignorException("No suitable new owner was found for the partition" + partition);
+ }
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromTargetAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToTargetAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberPair(" +
+ "srcMemberId='" + srcMemberId + '\'' +
+ ", dstMemberId='" + dstMemberId + '\'' +
+ ')';
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among members.
+ * During a partition rebalance, it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(
+ TopicIdPartition partition,
+ String oldOwner,
+ String newOwner
+ ) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ String expectedOldOwner = partitionMovementsByPartition.get(partition).dstMemberId;
+ if (!oldOwner.equals(expectedOldOwner)) {
+ throw new PartitionAssignorException("Old owner does not match expected value for partition: " + partition);
+ }
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
Review Comment:
oya thanks for the observation, I actually kept it because the sizes data corresponds to that of a particular member, if we want to print it or check equality it's better to have the memberId. What do you think?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "jeffkbkim (via GitHub)" <gi...@apache.org>.
jeffkbkim commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396498916
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,828 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * List of all members sorted by their assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that are retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Maintains a sorted set of consumers based on how many topic partitions are already assigned to them.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client side.
+ *
+ * Only populated with partitions that weren't retained due to a rack mismatch when rack aware strategy is used.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(allTopicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
Review Comment:
i think what you have works. it's simple and it fits in a single line.
an alternative would be
```
members.keySet().forEach(memberId ->
targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()))
);
```
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396612728
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
Review Comment:
I just wanted to confirm, by those two did you mean getting the subscibed topic ids?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396607932
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
Review Comment:
makes sense, and since we only use it to iterate once anyways, I've changed it to just iterate through it directly.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396625136
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
+
+ modified = false;
+ reassignmentOccurred = false;
+ // Reassign all reassignable partitions sorted in descending order
+ // by totalPartitions/number of subscribed members,
+ // until the full list is processed or a balance is achieved.
+ List<TopicIdPartition> reassignablePartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition reassignablePartition : reassignablePartitions) {
+ // Only check if there is any change in balance if any moves were made.
+ if (reassignmentOccurred && isBalanced()) {
+ break;
+ }
+ reassignmentOccurred = false;
+
+ // The topicIdPartition must have at least two consumers.
+ if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
+ LOG.error("Expected more than one potential member for topicIdPartition '{}'", reassignablePartition);
+
+ // The topicIdPartition must have a current member.
+ String member = partitionOwnerInTargetAssignment.get(reassignablePartition);
+ if (member == null)
+ LOG.error("Expected topicIdPartition '{}' to be assigned to a member", reassignablePartition);
+
+ boolean foundMatchingRackMember = false;
+
+ // If rack strategy is used and the current assignment adheres to rack matching,
+ // check if another member in the same rack is better suited for this topicIdPartition.
+ if (rackInfo.useRackStrategy) {
+
+ String memberRack = rackInfo.memberRacks.get(member);
+ Set<String> partitionRacks = rackInfo.partitionRacks.get(reassignablePartition);
+
+ if (partitionRacks.contains(memberRack)) {
+ for (String otherMember : rackInfo.getSortedMembersWithMatchingRack(reassignablePartition, targetAssignment)) {
+ // Only subscribed members eligible for re-balancing should be considered.
+ if (!sortedMembersByAssignmentSize.contains(otherMember) ||
+ !membersPerTopic.containsKey(reassignablePartition.topicId())
+ ) continue;
+
+ String otherMemberRack = rackInfo.memberRacks.get(otherMember);
+ if (otherMemberRack == null || !partitionRacks.contains(otherMemberRack))
+ continue;
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition, otherMember);
+ modified = true;
+ reassignmentOccurred = true;
+ foundMatchingRackMember = true;
+ break;
+ }
+ }
+ }
+ }
+
+ // If rack-aware strategy is not used OR no other member with matching rack was better suited,
+ // First check if the topicIdPartition already belongs to its previous owner, if not and a previous owner
+ // exists, check if the topicIdPartition can be assigned to it.
+ // If not, assign to any other better suited member with the topic subscription.
+ if (!foundMatchingRackMember) {
+ boolean isPartitionSticky = assignedStickyPartitions.contains(reassignablePartition);
+ boolean isCurrentOwnerKnown = currentPartitionOwners.containsKey(reassignablePartition);
+
+ if (rackInfo.useRackStrategy && !isPartitionSticky && isCurrentOwnerKnown) {
+ String currentOwner = currentPartitionOwners.get(reassignablePartition);
+ int currentMemberAssignmentSize = assignmentManager.targetAssignmentSize(member);
+ int currentOwnerAssignmentSize = assignmentManager.targetAssignmentSize(currentOwner);
+
+ // Check if reassignment is needed based on assignment sizes
+ if (currentMemberAssignmentSize > currentOwnerAssignmentSize + 1) {
+ reassignPartition(reassignablePartition, currentOwner);
+ modified = true;
+ reassignmentOccurred = true;
+ }
+ }
+
+ // The topicIdPartition is already sticky and no other member with matching rack is better suited.
+ if (!reassignmentOccurred) {
+ for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
+ if (assignmentManager.targetAssignmentSize(member) >
+ assignmentManager.targetAssignmentSize(otherMember) + 1
+ ) {
+ reassignPartition(reassignablePartition);
+ modified = true;
+ reassignmentOccurred = true;
+ break;
+ }
+ }
+ }
+ }
+ }
+ } while (modified);
+ }
+
+ /**
+ * Reassigns a partition to an eligible member with the fewest current target assignments.
+ * <ul>
+ * <li> Iterates over members sorted by ascending assignment size. </li>
+ * <li> Selects the first member subscribed to the partition's topic. </li>
+ * </ul>
+ *
+ * @param partition The partition to reassign.
+ * @throws AssertionError If no subscribed member is found.
+ */
+ private void reassignPartition(TopicIdPartition partition) {
+ // Find the new member with the least assignment size.
+ String newOwner = null;
+ for (String anotherMember: sortedMembersByAssignmentSize) {
+ if (members.get(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
+ newOwner = anotherMember;
+ break;
+ }
+ }
+
+ assert newOwner != null;
+
+ reassignPartition(partition, newOwner);
+ }
+
+ /**
+ * Reassigns the given partition to a new member while considering partition movements and stickiness.
+ * <p>
+ * This method performs the following actions:
+ * <ol>
+ * <li> Determines the current owner of the partition. </li>
+ * <li> Identifies the correct partition to move, adhering to stickiness constraints. </li>
+ * <li> Processes the partition movement to the new member. </li>
+ * </ol>
+ *
+ * @param partition The {@link TopicIdPartition} to be reassigned.
+ * @param newMember The Id of the member to which the partition should be reassigned.
+ */
+ private void reassignPartition(TopicIdPartition partition, String newMember) {
+ String member = partitionOwnerInTargetAssignment.get(partition);
+ // Find the correct partition movement considering the stickiness requirement.
+ TopicIdPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(
+ partition,
+ member,
+ newMember
+ );
+ processPartitionMovement(partitionToBeMoved, newMember);
+ }
+
+ private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
+ String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+
+ partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+
+ assignmentManager.removePartitionFromMembersAssignment(topicIdPartition, oldMember);
+ assignmentManager.addPartitionToMembersAssignment(topicIdPartition, newMember);
+ }
+
+ /**
+ * <code> MemberPair </code> represents a pair of member Ids involved in a partition reassignment. Each
+ * <code> MemberPair </code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+ * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+ * partition of the particular topic was moved from the source member to the destination member during the rebalance.
+ */
+ private static class MemberPair {
+ private final String srcMemberId;
+ private final String dstMemberId;
+
+ MemberPair(String srcMemberId, String dstMemberId) {
+ this.srcMemberId = srcMemberId;
+ this.dstMemberId = dstMemberId;
+ }
+
+ public String toString() {
+ return this.srcMemberId + "->" + this.dstMemberId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+ result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (obj == null)
+ return false;
+
+ if (!getClass().isInstance(obj))
+ return false;
+
+ MemberPair otherPair = (MemberPair) obj;
+ return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+ }
+ }
+
+ /**
+ * This class maintains some data structures to simplify lookup of partition movements among consumers.
+ * At each point of time during a partition rebalance it keeps track of partition movements corresponding to each topic,
+ * and also possible movement (in form a <code>MemberPair</code> object) for each partition.
+ */
+ private static class PartitionMovements {
+ private final Map<Uuid, Map<MemberPair, Set<TopicIdPartition>>> partitionMovementsByTopic = new HashMap<>();
+ private final Map<TopicIdPartition, MemberPair> partitionMovementsByPartition = new HashMap<>();
+
+ private MemberPair removeMovementRecordOfPartition(TopicIdPartition partition) {
+ MemberPair pair = partitionMovementsByPartition.remove(partition);
+
+ Uuid topic = partition.topicId();
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ partitionMovementsForThisTopic.get(pair).remove(partition);
+ if (partitionMovementsForThisTopic.get(pair).isEmpty())
+ partitionMovementsForThisTopic.remove(pair);
+ if (partitionMovementsByTopic.get(topic).isEmpty())
+ partitionMovementsByTopic.remove(topic);
+
+ return pair;
+ }
+
+ private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
+ partitionMovementsByPartition.put(partition, pair);
+
+ Uuid topic = partition.topicId();
+ if (!partitionMovementsByTopic.containsKey(topic))
+ partitionMovementsByTopic.put(topic, new HashMap<>());
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ if (!partitionMovementsForThisTopic.containsKey(pair))
+ partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+ partitionMovementsForThisTopic.get(pair).add(partition);
+ }
+
+ private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
+ MemberPair pair = new MemberPair(oldOwner, newOwner);
+
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ // This partition was previously moved.
+ MemberPair existingPair = removeMovementRecordOfPartition(partition);
+ assert existingPair.dstMemberId.equals(oldOwner);
+ if (!existingPair.srcMemberId.equals(newOwner)) {
+ // The partition is not moving back to its previous consumer.
+ addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
+ }
+ } else
+ addPartitionMovementRecord(partition, pair);
+ }
+
+ private TopicIdPartition getTheActualPartitionToBeMoved(TopicIdPartition partition, String oldOwner, String newOwner) {
+ Uuid topic = partition.topicId();
+
+ if (!partitionMovementsByTopic.containsKey(topic))
+ return partition;
+
+ // This partition was previously moved.
+ if (partitionMovementsByPartition.containsKey(partition)) {
+ assert oldOwner.equals(partitionMovementsByPartition.get(partition).dstMemberId);
+ oldOwner = partitionMovementsByPartition.get(partition).srcMemberId;
+ }
+
+ Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+ MemberPair reversePair = new MemberPair(newOwner, oldOwner);
+ if (!partitionMovementsForThisTopic.containsKey(reversePair))
+ return partition;
+
+ return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+ }
+ }
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ protected class AssignmentManager {
+ private final Map<String, MemberAssignmentData> membersWithAssignmentSizes = new HashMap<>();
+
+ /**
+ * Represents the assignment metadata for a member.
+ */
+ private class MemberAssignmentData {
+ final String memberId;
+ int currentAssignmentSize = 0;
+ int maxAssignmentSize;
+
+ /**
+ * Constructs a MemberAssignmentData with the given member Id.
+ *
+ * @param memberId The Id of the member.
+ */
+ MemberAssignmentData(String memberId) {
+ this.memberId = memberId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ MemberAssignmentData that = (MemberAssignmentData) o;
+ return memberId.equals(that.memberId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(memberId);
+ }
+
+ @Override
+ public String toString() {
+ return "MemberAssignmentData{" +
Review Comment:
ohhh makes sense okay, these were all autogenerated, maybe I'll try to change the settings :/
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "jeffkbkim (via GitHub)" <gi...@apache.org>.
jeffkbkim commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1400045359
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
Review Comment:
perhaps we can add a test to trigger one of the following?
```
// If member is not subscribed to the partition's topic, return false without assigning.
if (!members.get(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
return false;
}
// If the member's current assignment is already at max, return false without assigning.
if (isMemberAtMaxCapacity(memberId)) {
return false;
}
```
from maybeAssignPartitionToMember()
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1384454036
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,828 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * List of all members sorted by their assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that are retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Maintains a sorted set of consumers based on how many topic partitions are already assigned to them.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client side.
+ *
+ * Only populated with partitions that weren't retained due to a rack mismatch when rack aware strategy is used.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(allTopicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
Review Comment:
nit: Could you give me a suggestion for the formatting of this?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1396681253
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
Review Comment:
we need to get all the topicIdPartitions again to be considered for reassignment, did you have any other suggestions to avoid rebuilding the set?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "dajac (via GitHub)" <gi...@apache.org>.
dajac commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1397600750
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId: subscriptionIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member: members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty() && !isBalanced()) performReassignments();
+ }
+
+ private void performReassignments() {
+ boolean modified;
+ boolean reassignmentOccurred;
+ // Repeat reassignment until no partition can be moved to improve the balance.
+ do {
+ // Before re-starting the round of reassignments check if the assignment is already balanced.
+ if (isBalanced()) break;
Review Comment:
`isBalanced` seems quite heavy so it would be better to avoid this. I wonder if we could just add it to the condition of the do..while loop.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "dajac (via GitHub)" <gi...@apache.org>.
dajac commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1403102326
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
Review Comment:
Created https://issues.apache.org/jira/browse/KAFKA-14048.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1400980116
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,903 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscribedTopicIds;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment of the member.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private final PartitionMovements partitionMovements;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscribedTopicIds = new HashSet<>();
+ this.membersPerTopic = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId -> {
+ // Check if the subscribed topic exists.
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ subscribedTopicIds.add(topicId);
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId);
+ targetAssignment.put(memberId, new MemberAssignment(new HashMap<>()));
+ });
+ });
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscribedTopicIds);
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.partitionMovements = new PartitionMovements();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscribedTopicIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId)) {
+ break;
+ }
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId())
+ );
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId : assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member : sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId : members.get(member).subscribedTopicIds()) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment after the initial round of assignments have completed.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty())
+ throw new PartitionAssignorException("Some partitions were left unassigned");
+ // Refill unassigned partitions with all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber));
+
+ // Narrow down the reassignment scope to only those partitions that can actually be reassigned.
+ Set<TopicIdPartition> fixedPartitions = new HashSet<>();
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ fixedPartitions.add(new TopicIdPartition(topicId, i));
+ }
+ }
+ }
+ unassignedPartitions.removeAll(fixedPartitions);
+
+ // Narrow down the reassignment scope to only those members that are subject to reassignment.
+ for (String member : members.keySet()) {
+ if (!canMemberParticipateInReassignment(member)) {
+ sortedMembersByAssignmentSize.remove(member);
+ }
+ }
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!unassignedPartitions.isEmpty()) performReassignments();
+ }
+
+ private void performReassignments() {
Review Comment:
hey I was worried about this too, I've thought of multiple ways to make it better, hence it took so long to implement, but seems like I've run out of ideas and any idea I had came with its own problems. I'm open to discussing it again.
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1401419710
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AbstractUniformAssignmentBuilder.java:
##########
@@ -223,22 +223,13 @@ protected boolean racksMismatch(String memberId, TopicIdPartition tp) {
* @return A sorted list of partitions with potential members in the same rack.
*/
protected List<TopicIdPartition> sortPartitionsByRackMembers(Collection<TopicIdPartition> topicIdPartitions) {
-
- List<TopicIdPartition> sortedPartitionsList = topicIdPartitions.stream()
- .sorted(Comparator.comparing(TopicIdPartition::topicId).thenComparing(TopicIdPartition::partitionId))
+ return topicIdPartitions.stream()
+ .filter(tp -> membersWithSameRackAsPartition.containsKey(tp) && !membersWithSameRackAsPartition.get(tp).isEmpty())
+ .sorted(Comparator.comparing(
+ (TopicIdPartition tp) -> membersWithSameRackAsPartition.getOrDefault(tp, Collections.emptyList()).size())
Review Comment:
If I omit the type then I get an error cause the topicId method can't be inferred, I need to explicitly mention the type. Do you have a workaround for this?
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1400884026
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
Review Comment:
So the two conditions in maybeAssignPartitionsToMember are a) is the member subscribed to the topic in question b) is the member at max capacity. The first condition cannot be true since we check for topic subscriptions before assigning sticky partitions. The second condition cannot be true either since max assignment refers to the sum total of all the subscribed partitions, which means all the sticky partitions would be assigned to the member anyways. I'm going to change it to directly assign partition to member in this case
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-15484: General Rack Aware Assignor [kafka]
Posted by "rreddy-22 (via GitHub)" <gi...@apache.org>.
rreddy-22 commented on code in PR #14481:
URL: https://github.com/apache/kafka/pull/14481#discussion_r1399989044
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/GeneralUniformAssignmentBuilder.java:
##########
@@ -14,17 +14,891 @@
* 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.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+/**
+ * The general uniform assignment builder is used to generate the target assignment for a consumer group with
+ * at least one of its members subscribed to a different set of topics.
+ *
+ * Assignments are done according to the following principles:
+ *
+ * <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 Matching: 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>
+ *
+ * This assignment builder prioritizes the above properties in the following order:
+ * Balance > Rack Matching > Stickiness.
+ */
public class GeneralUniformAssignmentBuilder extends AbstractUniformAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(GeneralUniformAssignmentBuilder.class);
+ /**
+ * The member metadata obtained from the assignment specification.
+ */
+ private final Map<String, AssignmentMemberSpec> members;
+
+ /**
+ * The topic and partition metadata describer.
+ */
+ private final SubscribedTopicDescriber subscribedTopicDescriber;
+
+ /**
+ * The list of all the topic Ids that the consumer group is subscribed to.
+ */
+ private final Set<Uuid> subscriptionIds;
+
+ /**
+ * Rack information.
+ */
+ private final RackInfo rackInfo;
+
+ /**
+ * List of subscribed members for each topic.
+ */
+ private final Map<Uuid, List<String>> membersPerTopic;
+
+ /**
+ * The partitions that still need to be assigned.
+ */
+ private final Set<TopicIdPartition> unassignedPartitions;
+
+ /**
+ * All the partitions that have been retained from the existing assignment.
+ */
+ private final Set<TopicIdPartition> assignedStickyPartitions;
+
+ /**
+ * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ */
+ private final AssignmentManager assignmentManager;
+
+ /**
+ * List of all the members sorted by their respective assignment sizes.
+ */
+ private final TreeSet<String> sortedMembersByAssignmentSize;
+
+ /**
+ * Tracks the owner of each partition in the existing assignment on the client.
+ *
+ * Only populated when rack aware strategy is used.
+ * Contains partitions that weren't retained due to a rack mismatch.
+ */
+ private final Map<TopicIdPartition, String> currentPartitionOwners;
+
+ /**
+ * Tracks the owner of each partition in the target assignment.
+ */
+ private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+
+ /**
+ * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ */
+ private PartitionMovements partitionMovements;
+
+ /**
+ * The new assignment that will be returned.
+ */
+ private final Map<String, MemberAssignment> targetAssignment;
+
+ public GeneralUniformAssignmentBuilder(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
+ this.members = assignmentSpec.members();
+ this.subscribedTopicDescriber = subscribedTopicDescriber;
+ this.subscriptionIds = assignmentSpec.members().values().stream()
+ .flatMap(memberSpec -> memberSpec.subscribedTopicIds().stream())
+ .peek(topicId -> {
+ int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
+ if (partitionCount == -1) {
+ throw new PartitionAssignorException(
+ "Members are subscribed to topic " + topicId + " which doesn't exist in the topic metadata."
+ );
+ }
+ })
+ .collect(Collectors.toSet());
+ this.rackInfo = new RackInfo(assignmentSpec, subscribedTopicDescriber, subscriptionIds);
+ this.membersPerTopic = new HashMap<>();
+ members.forEach((memberId, memberMetadata) -> {
+ Collection<Uuid> topics = memberMetadata.subscribedTopicIds();
+ topics.forEach(topicId ->
+ membersPerTopic.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberId)
+ );
+ });
+ this.unassignedPartitions = new HashSet<>(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
+ this.assignedStickyPartitions = new HashSet<>();
+ this.assignmentManager = new AssignmentManager();
+ this.sortedMembersByAssignmentSize = assignmentManager.getSortedMembersByAssignmentSize(members.keySet());
+ this.currentPartitionOwners = new HashMap<>();
+ this.partitionOwnerInTargetAssignment = new HashMap<>();
+ this.targetAssignment = new HashMap<>();
+ }
+
+ /**
+ * Here's the step-by-step breakdown of the assignment process:
+ *
+ * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
+ * <ul><li> If a partition's rack mismatches with its owner, track it for future use. </li></ul>
+ * <li> If rack aware strategy is possible, allocate unassigned partitions to members in the same rack. </li>
+ * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner. If possible, allocate
+ * the partition back to it's existing owner in case it was not retained earlier due to a rack mismatch. </li>
+ * <li> Iterate through the assignment until it is balanced. </li>
+ */
@Override
protected GroupAssignment buildAssignment() {
- return null;
+ if (subscriptionIds.isEmpty()) {
+ LOG.info("The subscription list is empty, returning an empty assignment");
+ return new GroupAssignment(Collections.emptyMap());
+ }
+
+ members.keySet().forEach(memberId -> targetAssignment.put(memberId, new MemberAssignment(new HashMap<>())));
+ partitionMovements = new PartitionMovements();
+
+ // When rack awareness is enabled, only sticky partitions with matching rack are retained.
+ // Otherwise, all existing partitions are retained until max assignment size.
+ assignStickyPartitions();
+
+ if (rackInfo.useRackStrategy) rackAwarePartitionAssignment();
+ unassignedPartitionsAssignment();
+
+ balance();
+
+ return new GroupAssignment(targetAssignment);
+ }
+
+ /**
+ * <li> TopicIdPartitions are sorted in descending order based on the value:
+ * totalPartitions/number of subscribed members. </li>
+ * <li> If the above value is the same then topicIdPartitions are sorted in
+ * ascending order of number of subscribers. </li>
+ * <li> If both criteria are the same, sort in ascending order of the partition Id.
+ * This last criteria is for predictability of the assignments. </li>
+ *
+ * @param topicIdPartitions The topic partitions that need to be sorted.
+ * @return A list of sorted topic partitions.
+ */
+ private List<TopicIdPartition> sortTopicIdPartitions(Collection<TopicIdPartition> topicIdPartitions) {
+ Comparator<TopicIdPartition> comparator = Comparator
+ .comparingDouble((TopicIdPartition topicIdPartition) -> {
+ int totalPartitions = subscribedTopicDescriber.numPartitions(topicIdPartition.topicId());
+ int totalSubscribers = membersPerTopic.get(topicIdPartition.topicId()).size();
+ return (double) totalPartitions / totalSubscribers;
+ })
+ .reversed()
+ .thenComparingInt(topicIdPartition -> membersPerTopic.get(topicIdPartition.topicId()).size())
+ .thenComparingInt(TopicIdPartition::partitionId);
+
+ return topicIdPartitions.stream()
+ .sorted(comparator)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Gets a set of partitions that are to be retained from the existing assignment. This includes:
+ * <li> Partitions from topics that are still present in both the new subscriptions and the topic metadata. </li>
+ * <li> When using a rack-aware strategy, only partitions with member owners in the same rack are retained. </li>
+ * <li> Track current partition owners when there is a rack mismatch. </li>
+ */
+ private void assignStickyPartitions() {
+ members.forEach((memberId, assignmentMemberSpec) ->
+ assignmentMemberSpec.assignedPartitions().forEach((topicId, currentAssignment) -> {
+ if (assignmentMemberSpec.subscribedTopicIds().contains(topicId)) {
+ currentAssignment.forEach(partition -> {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
+ if (rackInfo.useRackStrategy && rackInfo.racksMismatch(memberId, topicIdPartition)) {
+ currentPartitionOwners.put(topicIdPartition, memberId);
+ } else {
+ assignmentManager.maybeAssignPartitionToMember(topicIdPartition, memberId);
+ assignedStickyPartitions.add(topicIdPartition);
+ }
+ });
+ } else {
+ LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
+ }
+ })
+ );
+ }
+
+ /**
+ * Allocates the unassigned partitions to members in the same rack, if available.
+ */
+ private void rackAwarePartitionAssignment() {
+ // Sort partitions in ascending order by the number of potential members with matching racks.
+ // Only partitions with potential members in the same rack are returned.
+ List<TopicIdPartition> sortedPartitions = rackInfo.sortPartitionsByRackMembers(unassignedPartitions);
+
+ sortedPartitions.forEach(partition -> {
+ List<String> sortedMembersByAssignmentSize = rackInfo.getSortedMembersWithMatchingRack(
+ partition,
+ targetAssignment
+ );
+
+ for (String memberId : sortedMembersByAssignmentSize) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, memberId))
+ break;
+ }
+ });
+ }
+
+ /**
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
+ * <li> Partitions are sorted to maximize the probability of a balanced assignment. </li>
+ * <li> If there was an assignment that wasn't retained due to a rack mismatch,
+ * check if the partition can retain its existing assignment. </li>
+ * <li> Sort members in ascending order of their current target assignment sizes
+ * to ensure the least filled member gets the partition first. </li>
+ */
+ private void unassignedPartitionsAssignment() {
+ List<TopicIdPartition> sortedPartitions = sortTopicIdPartitions(unassignedPartitions);
+
+ for (TopicIdPartition partition : sortedPartitions) {
+ if (rackInfo.useRackStrategy && currentPartitionOwners.containsKey(partition)) {
+ String prevOwner = currentPartitionOwners.get(partition);
+ if (assignmentManager.maybeAssignPartitionToMember(partition, prevOwner)) {
+ continue;
+ }
+ }
+
+ TreeSet<String> sortedMembers = assignmentManager.getSortedMembersByAssignmentSize(
+ membersPerTopic.get(partition.topicId()));
+
+ for (String member : sortedMembers) {
+ if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * If a topic has two or more potential members it is subject to reassignment.
+ *
+ * @return true if the topic can participate in reassignment, false otherwise.
+ */
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return membersPerTopic.get(topicId).size() >= 2;
+ }
+
+ /**
+ * If a member is not assigned all its potential partitions it is subject to reassignment.
+ * If any of the partitions assigned to a member is subject to reassignment, the member itself
+ * is subject to reassignment.
+ *
+ * @return true if the member can participate in reassignment, false otherwise.
+ */
+ private boolean canMemberParticipateInReassignment(String memberId) {
+ Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).targetPartitions().keySet();
+
+ int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+
+ if (currentAssignmentSize > maxAssignmentSize)
+ LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+
+ if (currentAssignmentSize < maxAssignmentSize)
+ return true;
+
+ for (Uuid topicId: assignedTopicIds) {
+ if (canTopicParticipateInReassignment(topicId))
+ return true;
+ }
+ return false;
+ }
+
+ /**
+ * Determine if the current assignment is a balanced one.
+ *
+ * @return true if the given assignment is balanced; false otherwise.
+ */
+ private boolean isBalanced() {
+ int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
+ int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
+
+ // If minimum and maximum numbers of partitions assigned to consumers differ by at most one return true.
+ if (min >= max - 1)
+ return true;
+
+ // For each member that does not have all the topic partitions it can get make sure none of the
+ // topic partitions it could but did not get cannot be moved to it (because that would break the balance).
+ // Members with the least assignment sizes are checked first to see if they can receive any more partitions.
+ for (String member: sortedMembersByAssignmentSize) {
+ int memberPartitionCount = assignmentManager.targetAssignmentSize(member);
+
+ // Skip if this member already has all the topic partitions it can get.
+ List<Uuid> allSubscribedTopics = new ArrayList<>(members.get(member).subscribedTopicIds());
+ int maxAssignmentSize = assignmentManager.maxAssignmentSize(member);
+
+ if (memberPartitionCount == maxAssignmentSize)
+ continue;
+
+ // Otherwise make sure it cannot get any more partitions.
+ for (Uuid topicId: allSubscribedTopics) {
+ Set<Integer> assignedPartitions = targetAssignment.get(member).targetPartitions().get(topicId);
+ for (int i = 0; i < subscribedTopicDescriber.numPartitions(topicId); i++) {
+ TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, i);
+ if (assignedPartitions == null || !assignedPartitions.contains(i)) {
+ String otherMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ int otherMemberPartitionCount = assignmentManager.targetAssignmentSize(otherMember);
+ if (memberPartitionCount + 1 < otherMemberPartitionCount) {
+ LOG.debug("{} can be moved from member {} to member {} for a more balanced assignment.",
+ topicIdPartition, otherMember, member);
+ return false;
+ }
+ }
+ }
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Balance the current assignment using the data structures created in the assignPartitions(...) method above.
+ */
+ private void balance() {
+ if (!unassignedPartitions.isEmpty()) unassignedPartitionsAssignment();
+ // Refill unassigned partitions will all the topicId partitions.
+ unassignedPartitions.addAll(topicIdPartitions(subscriptionIds, subscribedTopicDescriber));
Review Comment:
discussed on zoom
--
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: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org