You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by da...@apache.org on 2024/10/17 08:43:35 UTC
(kafka) branch trunk updated: KAFKA-17710; Rework uniform heterogeneous assignor to improve perf (#17385)
This is an automated email from the ASF dual-hosted git repository.
dajac pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new 017da210999 KAFKA-17710; Rework uniform heterogeneous assignor to improve perf (#17385)
017da210999 is described below
commit 017da210999c72789ce3b720af04c1a834b80a5c
Author: Sean Quah <sq...@confluent.io>
AuthorDate: Thu Oct 17 09:43:29 2024 +0100
KAFKA-17710; Rework uniform heterogeneous assignor to improve perf (#17385)
Rework the uniform heterogeneous assignor to improve performance, while
preserving the high level ideas and structure from the existing
implementation:
* The assignor works in 3 stages: importing the previous assignment for
stickiness, assigning unassigned partitions and iteratively
reassigning partitions to improve balance.
* Unassigned partitions are assigned to the subscribers with the least
number of partitions. This maximizes balance within a single topic.
* During the iterative rebalancing phase, partitions are reassigned to
their previous owner if it improves balance (stickiness restoration).
* During the iterative rebalancing phase, partitions are reassigned to
the subscriber with the least number of partitions to improve
balance.
A non-exhaustive list of changes is:
* The assignment of unassigned partitions and iterative reassignment
stages now works through partitions topic by topic. Previously
partitions from topics with the same number of partitions per
subscriber would be interleaved. Since we iterate topic by topic, we
can reuse data about topic subscribers.
* Instead of maintaining TreeSets to find the least loaded subscribers,
we sort an ArrayList of subscribers once per topic and start filling
up subscribers, least loaded first. In testing, this approach was
found to be faster than maintaining PriorityQueues.
* Implement stickiness restoration by creating a mapping of partitions
to previous owner and checking against that mapping, instead of
tracking partition movements during iterative reassignment.
* Track member partition counts using a plain int array, to avoid
overhead from boxing and HashMap lookups. Member partition counts are
accessed very frequently and this needs to be fast. As a consequence,
we have to number members 0 to M - 1.
* Bound the iterative reassignment stage to a fixed number of
iterations. Under some uncommon subscription patterns, the iterative
reassignment stage converges slowly. In these cases, the iterative
reassignment stage terminates without producing an optimally balanced
assignment anyway (see javadoc for balanceTopics).
* Re-use Maps from the previous assignment where possible,
ie. introduce a copy-on-write mechanism while computing the new
assignment.
Reviewers: David Jacot <dj...@confluent.io>
---
.../group/assignor/AssignorHelpers.java | 54 +
.../UniformHeterogeneousAssignmentBuilder.java | 1326 ++++++++++----------
.../UniformHomogeneousAssignmentBuilder.java | 29 +-
.../UniformHeterogeneousAssignmentBuilderTest.java | 134 +-
4 files changed, 867 insertions(+), 676 deletions(-)
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignorHelpers.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignorHelpers.java
new file mode 100644
index 00000000000..2b30b0b3ae9
--- /dev/null
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/AssignorHelpers.java
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group.assignor;
+
+import org.apache.kafka.common.Uuid;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Provides helper methods for assignors.
+ */
+public final class AssignorHelpers {
+ private static final Class<?> UNMODIFIABLE_MAP_CLASS = Collections.unmodifiableMap(new HashMap<>()).getClass();
+ private static final Class<?> EMPTY_MAP_CLASS = Collections.emptyMap().getClass();
+
+ /**
+ * @return True if the provided map is an UnmodifiableMap or EmptyMap. Those classes are not
+ * public hence we cannot use the `instanceof` operator.
+ */
+ public static boolean isImmutableMap(Map<?, ?> map) {
+ return UNMODIFIABLE_MAP_CLASS.isInstance(map) || EMPTY_MAP_CLASS.isInstance(map);
+ }
+
+ /**
+ * Deep copies a member assignment map.
+ * @param map The assignment to copy.
+ * @return A deep copy of the assignment.
+ */
+ public static Map<Uuid, Set<Integer>> deepCopyAssignment(Map<Uuid, Set<Integer>> map) {
+ Map<Uuid, Set<Integer>> copy = new HashMap<>(map.size());
+ for (Map.Entry<Uuid, Set<Integer>> entry : map.entrySet()) {
+ copy.put(entry.getKey(), new HashSet<>(entry.getValue()));
+ }
+ return copy;
+ }
+}
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilder.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilder.java
index 6760d1e8663..ca5ba77fbbd 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilder.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilder.java
@@ -23,12 +23,12 @@ import org.apache.kafka.coordinator.group.api.assignor.MemberAssignment;
import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignorException;
import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
import org.apache.kafka.coordinator.group.modern.MemberAssignmentImpl;
-import org.apache.kafka.server.common.TopicIdPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
@@ -36,22 +36,20 @@ 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 heterogeneous 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.
- *
+ * <p/>
* 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> Stickiness: Minimize partition movements among members by retaining
- * as much of the existing assignment as possible. </li>
+ * <ol>
+ * <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>Stickiness: Minimize partition movements among members by retaining
+ * as much of the existing assignment as possible.</li>
+ * </ol>
*
* This assignment builder prioritizes the above properties in the following order:
* Balance > Stickiness.
@@ -59,6 +57,19 @@ import java.util.stream.Collectors;
public class UniformHeterogeneousAssignmentBuilder {
private static final Logger LOG = LoggerFactory.getLogger(UniformHeterogeneousAssignmentBuilder.class);
+ /**
+ * The maximum number of iterations to perform in the final iterative balancing phase.
+ * <p/>
+ * When all subscribers to a topic have the same subscriptions, the final balancing phase
+ * converges in a single iteration and takes another iteration to detect that no further
+ * balancing is possible. This describes most subscription patterns.
+ * <p/>
+ * When subscribers have overlapping, non-identical subscriptions, the final balancing phase can
+ * take a lot longer to converge. We assume that each iteration roughly halves the imbalance and
+ * choose a limit that allows for a reasonable balance to be achieved without taking too long.
+ */
+ private static final int MAX_ITERATION_COUNT = 16;
+
/**
* The group metadata specification.
*/
@@ -75,89 +86,175 @@ public class UniformHeterogeneousAssignmentBuilder {
private final Set<Uuid> subscribedTopicIds;
/**
- * List of subscribed members for each topic.
+ * The list of members in the consumer group.
*/
- private final Map<Uuid, List<String>> membersPerTopic;
+ private final List<String> memberIds;
/**
- * The new assignment that will be returned.
+ * Maps member ids to their indices in the memberIds list.
*/
- private final Map<String, MemberAssignment> targetAssignment;
+ private final Map<String, Integer> memberIndices;
/**
- * The partitions that still need to be assigned.
+ * List of subscribed members for each topic, in ascending order.
+ * <p/>
+ * Members are stored as integer indices into the memberIds array.
*/
- private final Set<TopicIdPartition> unassignedPartitions;
+ private final Map<Uuid, List<Integer>> topicSubscribers;
/**
- * All the partitions that have been retained from the existing assignment.
+ * The new assignment that will be returned.
*/
- private final Set<TopicIdPartition> assignedStickyPartitions;
+ private final Map<String, MemberAssignment> targetAssignment;
/**
- * Manages assignments to members based on their current assignment size and maximum allowed assignment size.
+ * The number of partitions each member is assigned in the new assignment.
+ * <p/>
+ * The indices are the same as the memberIds array. memberTargetAssignmentSizes[i] contains the
+ * number of partitions assigned to memberIds[i].
+ * <p/>
+ * We use a plain int[] array here for performance reasons. Looking up the number of partitions
+ * assigned to a given member is a very common operation. By using a plain int[] array, we can
+ * avoid:
+ * <ul>
+ * <li>the overhead of boxing and unboxing integers.</li>
+ * <li>the cost of HashMap lookups.</li>
+ * <li>the cost of checking member ids for equality when there are HashMap collisions.</li>
+ * </ul>
*/
- private final AssignmentManager assignmentManager;
+ private final int[] memberTargetAssignmentSizes;
/**
- * List of all the members sorted by their respective assignment sizes.
+ * Orders topics by partitions per subscriber, descending.
+ * <p/>
+ * Ties are broken by subscriber count, ascending.
+ * <p/>
+ * Remaining ties are broken by topic id, ascending.
*/
- private final TreeSet<String> sortedMembersByAssignmentSize;
+ private final Comparator<Uuid> topicComparator;
/**
- * Tracks the owner of each partition in the target assignment.
+ * Orders members by their number of assigned partitions, ascending.
+ * <p/>
+ * Ties are broken by member index, ascending.
*/
- private final Map<TopicIdPartition, String> partitionOwnerInTargetAssignment;
+ private final Comparator<Integer> memberComparator;
/**
- * Handles all operations related to partition movements during a reassignment for balancing the target assignment.
+ * Tracks the owner of each partition in the target assignment.
+ * <p/>
+ * Owners are represented as indices into the memberIds array.
+ * -1 indicates an unassigned partition.
*/
- private final PartitionMovements partitionMovements;
+ private final Map<Uuid, int[]> targetAssignmentPartitionOwners;
public UniformHeterogeneousAssignmentBuilder(GroupSpec groupSpec, SubscribedTopicDescriber subscribedTopicDescriber) {
this.groupSpec = groupSpec;
this.subscribedTopicDescriber = subscribedTopicDescriber;
this.subscribedTopicIds = new HashSet<>();
- this.membersPerTopic = new HashMap<>();
+
+ // Number the members 0 to M - 1.
+ this.memberIds = new ArrayList<>(groupSpec.memberIds());
+ this.memberIndices = new HashMap<>();
+ for (int memberIndex = 0; memberIndex < this.memberIds.size(); memberIndex++) {
+ memberIndices.put(memberIds.get(memberIndex), memberIndex);
+ }
+
+ this.topicSubscribers = new HashMap<>();
+
this.targetAssignment = new HashMap<>();
- groupSpec.memberIds().forEach(memberId -> {
- groupSpec.memberSubscription(memberId).subscribedTopicIds().forEach(topicId -> {
+ this.memberTargetAssignmentSizes = new int[this.memberIds.size()];
+
+ // Build set of all subscribed topics and sets of subscribers per topic.
+ for (int memberIndex = 0; memberIndex < this.memberIds.size(); memberIndex++) {
+ String memberId = this.memberIds.get(memberIndex);
+ for (Uuid topicId : groupSpec.memberSubscription(memberId).subscribedTopicIds()) {
// Check if the subscribed topic exists.
- int partitionCount = subscribedTopicDescriber.numPartitions(topicId);
- if (partitionCount == -1) {
+ int numPartitions = subscribedTopicDescriber.numPartitions(topicId);
+ if (numPartitions == -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 MemberAssignmentImpl(new HashMap<>()));
- });
- this.unassignedPartitions = topicIdPartitions(subscribedTopicIds, subscribedTopicDescriber);
- this.assignedStickyPartitions = new HashSet<>();
- this.assignmentManager = new AssignmentManager(this.subscribedTopicDescriber);
- this.sortedMembersByAssignmentSize = assignmentManager.sortMembersByAssignmentSize(groupSpec.memberIds());
- this.partitionOwnerInTargetAssignment = new HashMap<>();
- this.partitionMovements = new PartitionMovements();
+ topicSubscribers.computeIfAbsent(topicId, k -> new ArrayList<>()).add(memberIndex);
+ }
+ }
+
+ this.topicComparator = new Comparator<Uuid>() {
+ @Override
+ public int compare(final Uuid topic1Id, final Uuid topic2Id) {
+ int topic1PartitionCount = subscribedTopicDescriber.numPartitions(topic1Id);
+ int topic2PartitionCount = subscribedTopicDescriber.numPartitions(topic2Id);
+ int topic1SubscriberCount = topicSubscribers.get(topic1Id).size();
+ int topic2SubscriberCount = topicSubscribers.get(topic2Id).size();
+
+ // Order by partitions per subscriber, descending.
+ int order = Double.compare(
+ (double) topic2PartitionCount / topic2SubscriberCount,
+ (double) topic1PartitionCount / topic1SubscriberCount
+ );
+
+ // Then order by subscriber count, ascending.
+ if (order == 0) {
+ order = Integer.compare(topic1SubscriberCount, topic2SubscriberCount);
+ }
+
+ // Then order by topic id, ascending.
+ if (order == 0) {
+ order = topic1Id.compareTo(topic2Id);
+ }
+
+ return order;
+ }
+ };
+
+ this.memberComparator = new Comparator<Integer>() {
+ @Override
+ public int compare(final Integer memberIndex1, final Integer memberIndex2) {
+ // Order by number of assigned partitions, ascending.
+ int order = Integer.compare(
+ memberTargetAssignmentSizes[memberIndex1],
+ memberTargetAssignmentSizes[memberIndex2]
+ );
+
+ // Then order by member index, ascending.
+ if (order == 0) {
+ order = memberIndex1.compareTo(memberIndex2);
+ }
+
+ return order;
+ }
+ };
+
+ // Initialize partition owners for the target assignments.
+ this.targetAssignmentPartitionOwners = new HashMap<>((int) ((this.subscribedTopicIds.size() / 0.75f) + 1));
+ for (Uuid topicId : this.subscribedTopicIds) {
+ int numPartitions = subscribedTopicDescriber.numPartitions(topicId);
+ int[] partitionOwners = new int[numPartitions];
+ Arrays.fill(partitionOwners, -1);
+ this.targetAssignmentPartitionOwners.put(topicId, partitionOwners);
+ }
}
/**
* Here's the step-by-step breakdown of the assignment process:
- *
- * <li> Retain partitions from the existing assignments a.k.a sticky partitions. </li>
- * <li> Allocate all the remaining unassigned partitions to the members in a balanced manner.</li>
- * <li> Iterate through the assignment until it is balanced. </li>
+ * <ol>
+ * <li>Revoke partitions from the existing assignment that are no longer part of each member's
+ * subscriptions.</li>
+ * <li>Allocate all the remaining unassigned partitions to the members in a balanced manner.</li>
+ * <li>Iterate through the assignment until it is balanced.</li>
+ * </ol>
*/
public GroupAssignment build() {
if (subscribedTopicIds.isEmpty()) {
- LOG.info("The subscription list is empty, returning an empty assignment");
return new GroupAssignment(Collections.emptyMap());
}
- assignStickyPartitions();
+ maybeRevokePartitions();
- unassignedPartitionsAssignment();
+ Map<Uuid, List<Integer>> unassignedPartitions = computeUnassignedPartitions();
+ assignRemainingPartitions(unassignedPartitions);
balance();
@@ -165,708 +262,665 @@ public class UniformHeterogeneousAssignmentBuilder {
}
/**
- * <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>
+ * Sorts topics based on the following criteria:
+ * <ol>
+ * <li>Topics are sorted in descending order of totalPartitions / number of subscribers
+ * (partitions per subscriber).</li>
+ * <li>Ties are broken by sorting in ascending order of number of subscribers.</li>
+ * <li>Any remaining ties are broken by sorting in ascending order of topic id.</li>
+ * </ol>
*
- * @param topicIdPartitions The topic partitions that need to be sorted.
- * @return A list of sorted topic partitions.
+ * The last criteria is for predictability of assignments.
+ *
+ * @param topicIds The topic ids that need to be sorted.
+ * @return A list of sorted topic ids.
*/
- 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());
+ private List<Uuid> sortTopicIds(Collection<Uuid> topicIds) {
+ List<Uuid> sortedTopicIds = new ArrayList<>(topicIds);
+ sortedTopicIds.sort(topicComparator);
+ return sortedTopicIds;
}
/**
- * 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>
+ * Revoke the partitions that are not part of each member's subscriptions.
+ *
+ * This method ensures that the original assignment is not copied if it is not altered.
*/
- private void assignStickyPartitions() {
- groupSpec.memberIds().forEach(memberId ->
- groupSpec.memberAssignment(memberId).partitions().forEach((topicId, currentAssignment) -> {
- if (groupSpec.memberSubscription(memberId).subscribedTopicIds().contains(topicId)) {
- currentAssignment.forEach(partition -> {
- TopicIdPartition topicIdPartition = new TopicIdPartition(topicId, partition);
- assignmentManager.addPartitionToTargetAssignment(topicIdPartition, memberId);
- assignedStickyPartitions.add(topicIdPartition);
- });
- } else {
- LOG.debug("The topic " + topicId + " is no longer present in the subscribed topics list");
- }
- })
- );
- }
+ private void maybeRevokePartitions() {
+ for (String memberId : groupSpec.memberIds()) {
+ int memberIndex = memberIndices.get(memberId);
- /**
- * 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> 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);
+ Map<Uuid, Set<Integer>> oldAssignment = groupSpec.memberAssignment(memberId).partitions();
+ Map<Uuid, Set<Integer>> newAssignment = null;
- for (TopicIdPartition partition : sortedPartitions) {
- TreeSet<String> sortedMembers = assignmentManager.sortMembersByAssignmentSize(
- membersPerTopic.get(partition.topicId())
- );
+ // The assignor expects to receive the assignment as an immutable map. It leverages
+ // this knowledge in order to avoid having to copy all assignments.
+ if (!AssignorHelpers.isImmutableMap(oldAssignment)) {
+ throw new IllegalStateException("The assignor expect an immutable map.");
+ }
- for (String member : sortedMembers) {
- if (assignmentManager.maybeAssignPartitionToMember(partition, member)) {
- break;
+ for (Map.Entry<Uuid, Set<Integer>> topicPartitions : oldAssignment.entrySet()) {
+ Uuid topicId = topicPartitions.getKey();
+ Set<Integer> partitions = topicPartitions.getValue();
+
+ if (groupSpec.memberSubscription(memberId).subscribedTopicIds().contains(topicId)) {
+ memberTargetAssignmentSizes[memberIndex] += partitions.size();
+
+ for (int partition : partitions) {
+ targetAssignmentPartitionOwners.get(topicId)[partition] = memberIndex;
+ }
+ } else {
+ if (newAssignment == null) {
+ // If the new assignment is null, we create a deep copy of the
+ // original assignment so that we can alter it.
+ newAssignment = AssignorHelpers.deepCopyAssignment(oldAssignment);
+ }
+ // Remove the entire topic.
+ newAssignment.remove(topicId);
}
}
- }
- }
- /**
- * 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 (newAssignment == null) {
+ newAssignment = oldAssignment;
+ }
+ targetAssignment.put(memberId, new MemberAssignmentImpl(newAssignment));
+ }
}
/**
- * 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.
+ * A class that holds the state for the balancing process.
+ * <p/>
+ * Keeps track of two ranges of members: the members with the least number of assigned
+ * partitions and the members with the most number of assigned partitions. Within each range,
+ * all members start with the same number of assigned partitions.
*/
- private boolean canMemberParticipateInReassignment(String memberId) {
- Set<Uuid> assignedTopicIds = targetAssignment.get(memberId).partitions().keySet();
+ private final class MemberAssignmentBalancer {
+ private final int[] memberTargetAssignmentSizes;
- int currentAssignmentSize = assignmentManager.targetAssignmentSize(memberId);
- int maxAssignmentSize = assignmentManager.maxAssignmentSize(memberId);
+ /**
+ * The members to balance, sorted by number of assigned partitions.
+ * <p/>
+ * Can be visualized like so:
+ * <pre>
+ * ^
+ * | #
+ * partition | ###
+ * count | ###
+ * | ##########
+ * +----------->
+ * members
+ * </pre>
+ */
+ private final List<Integer> sortedMembers;
- if (currentAssignmentSize > maxAssignmentSize)
- LOG.error("The member {} is assigned more partitions than the maximum possible.", memberId);
+ /**
+ * [0, leastLoadedRangeEnd) represents the range of members with the smallest partition
+ * count.
+ */
+ private int leastLoadedRangeEnd = 0; // exclusive
- if (currentAssignmentSize < maxAssignmentSize)
- return true;
+ /**
+ * The partition count of each member in the least loaded range before we start assigning
+ * partitions to them.
+ */
+ private int leastLoadedRangePartitionCount = -1;
- for (Uuid topicId : assignedTopicIds) {
- if (canTopicParticipateInReassignment(topicId))
- return true;
- }
- return false;
- }
+ /**
+ * The index of the next element in sortedMembers to which to assign the next partition.
+ */
+ private int nextLeastLoadedMember = 0;
- /**
- * Checks if the current assignments of partitions to members is balanced.
- *
- * Balance is determined by first checking if the difference in the number of partitions assigned
- * to any two members is one. If this is not true, it verifies that no member can
- * receive additional partitions without disrupting the balance.
- *
- * @return true if the assignment is balanced, false otherwise.
- */
+ /**
+ * [mostLoadedRangeStart, mostLoadedRangeEnd) is the range of members with the largest
+ * partition count.
+ */
+ private int mostLoadedRangeStart = 0; // inclusive
- private boolean isBalanced() {
- int min = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.first());
- int max = assignmentManager.targetAssignmentSize(sortedMembersByAssignmentSize.last());
-
- // If minimum and maximum numbers of partitions assigned to members differ by at most one return true.
- if (min >= max - 1)
- return true;
-
- // Ensure that members without a complete set of topic partitions cannot receive any additional partitions.
- // This maintains balance. Start by checking members with the fewest assigned partitions to see if they can take more.
- 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 : groupSpec.memberSubscription(member).subscribedTopicIds()) {
- Set<Integer> assignedPartitions = targetAssignment.get(member).partitions().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;
- }
+ /**
+ * [mostLoadedRangeStart, mostLoadedRangeEnd) is the range of members with the largest
+ * partition count.
+ * <p/>
+ * mostLoadedRangeEnd is not simply the end of the member list, since we need to be able to
+ * exclude heavily loaded members which do not have any partitions to give up.
+ */
+ private int mostLoadedRangeEnd = 0; // inclusive
- /**
- * 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));
+ /**
+ * The partition count of each member in the most loaded range before we start unassigning
+ * partitions from them.
+ */
+ private int mostLoadedRangePartitionCount = 1;
- // 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));
- }
- }
+ /**
+ * The index of the next element in sortedMembers from which to unassign the next partition.
+ */
+ private int nextMostLoadedMember = -1;
+
+ // The ranges can be visualized like so:
+ //
+ // most
+ // loaded
+ // ^ least range
+ // | loaded #
+ // | range [##<- ]##
+ // partition | ##[#######]##
+ // count | [...-> ]##[#######]##
+ // | [#######]##[#######]##
+ // +----------------------->
+ // members
+ // ^^
+ // These members have no partitions to give up.
+ //
+ // The least loaded range expands rightwards and the most loaded range expands leftwards.
+ // Note that we are not done once the ranges touch, but rather when the partition counts in
+ // the ranges differ by one or less.
+
+ public MemberAssignmentBalancer() {
+ this.memberTargetAssignmentSizes = UniformHeterogeneousAssignmentBuilder.this.memberTargetAssignmentSizes;
+ this.sortedMembers = new ArrayList<>(memberTargetAssignmentSizes.length);
}
- unassignedPartitions.removeAll(fixedPartitions);
- // Narrow down the reassignment scope to only those members that are subject to reassignment.
- for (String member : groupSpec.memberIds()) {
- if (!canMemberParticipateInReassignment(member)) {
- sortedMembersByAssignmentSize.remove(member);
+ /**
+ * Resets the state of the balancer for a new set of members.
+ *
+ * @param members The members to balance.
+ * @return The difference in partition counts between the most and least loaded members.
+ */
+ public int initialize(List<Integer> members) {
+ if (members.isEmpty()) {
+ // nextLeastLoadedMember needs at least one member.
+ throw new IllegalArgumentException("Cannot balance an empty subscriber list.");
}
- }
- // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
- if (!unassignedPartitions.isEmpty()) performReassignments();
- }
+ // Build a sorted list of subscribers, where members with the smallest partition count
+ // appear first.
+ sortedMembers.clear();
+ sortedMembers.addAll(members);
+ sortedMembers.sort(memberComparator);
- /**
- * Performs reassignments of partitions to balance the load across members.
- * This method iteratively reassigns partitions until no further moves can improve the balance.
- *
- * The method uses a do-while loop to ensure at least one pass over the partitions and continues
- * reassigning as long as there are modifications to the current assignments. It checks for balance
- * after each reassignment and exits if the balance is achieved.
- *
- * @throws PartitionAssignorException if there are inconsistencies in expected members per partition
- * or if a partition is expected to already be assigned but isn't.
- */
- private void performReassignments() {
- boolean modified;
- boolean reassignmentOccurred;
- // Repeat reassignment until no partition can be moved to improve the balance.
- do {
- 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;
+ leastLoadedRangeEnd = 0; // exclusive
+ leastLoadedRangePartitionCount = memberTargetAssignmentSizes[sortedMembers.get(0)] - 1;
+ nextLeastLoadedMember = 0;
+
+ mostLoadedRangeStart = sortedMembers.size(); // inclusive
+ mostLoadedRangeEnd = sortedMembers.size(); // exclusive
+ mostLoadedRangePartitionCount = memberTargetAssignmentSizes[sortedMembers.get(sortedMembers.size() - 1)] + 1;
+ nextMostLoadedMember = sortedMembers.size() - 1;
+
+ return memberTargetAssignmentSizes[sortedMembers.get(sortedMembers.size() - 1)] -
+ memberTargetAssignmentSizes[sortedMembers.get(0)];
+ }
+
+ /**
+ * Gets the least loaded member to which to assign a partition.
+ * <p/>
+ * Advances the state of the balancer by assuming that an additional partition has been
+ * assigned to the returned member.
+ * <p/>
+ * Assumes that there is at least one member.
+ *
+ * @return The least loaded member to which to assign a partition
+ */
+ public int nextLeastLoadedMember() {
+ // To generate a balanced assignment, we assign partitions to members in the
+ // [0, leastLoadedRangeEnd) range first. Once each member in the range has received a
+ // partition, the partition count of the range rises by one and we can try to expand it.
+ // Range full, Range full,
+ // bump level by one bump level by one,
+ // can't expand range. expand range.
+ // ^ ^ ^
+ // | # | # | [..-> ]#
+ // partition | ### -> | [..-> ]### -> | [.......##]#
+ // count | [..-> ]### | [.......]### | [.......##]#
+ // | [#######]### | [#######]### | [#########]#
+ // +-------------> +-------------> +------------>
+ // members members members
+
+ if (nextLeastLoadedMember >= leastLoadedRangeEnd) {
+ // We've hit the end of the range. Bump its level and try to expand it.
+ leastLoadedRangePartitionCount++;
+
+ // Expand the range.
+ while (leastLoadedRangeEnd < sortedMembers.size() &&
+ memberTargetAssignmentSizes[sortedMembers.get(leastLoadedRangeEnd)] == leastLoadedRangePartitionCount) {
+ leastLoadedRangeEnd++;
}
- reassignmentOccurred = false;
- // The topicIdPartition must have at least two members.
- if (membersPerTopic.get(reassignablePartition.topicId()).size() <= 1)
- throw new PartitionAssignorException(String.format("Expected more than one potential member for " +
- "topicIdPartition '%s'", reassignablePartition)
- );
+ // Reset the pointer to the start of the range of members with identical
+ // partition counts.
+ nextLeastLoadedMember = 0;
+ }
- // The topicIdPartition must have a current target owner.
- String currentTargetOwner = partitionOwnerInTargetAssignment.get(reassignablePartition);
- if (currentTargetOwner == null)
- throw new PartitionAssignorException(String.format("Expected topicIdPartition '%s' to be assigned " +
- "to a member", reassignablePartition)
- );
+ int leastLoadedMemberIndex = sortedMembers.get(nextLeastLoadedMember);
+ nextLeastLoadedMember++;
- for (String otherMember : membersPerTopic.get(reassignablePartition.topicId())) {
- if (assignmentManager.targetAssignmentSize(currentTargetOwner) > assignmentManager.targetAssignmentSize(otherMember) + 1) {
- reassignPartition(reassignablePartition);
- modified = true;
- reassignmentOccurred = true;
- break;
- }
+ return leastLoadedMemberIndex;
+ }
+
+ /**
+ * Gets the most loaded member from which to reassign a partition.
+ * <p/>
+ * Advances the state of the balancer by assuming that a partition has been unassigned from
+ * the returned member.
+ *
+ * @return The most loaded member from which to reassign a partition, or -1 if no such
+ * member exists.
+ */
+ public int nextMostLoadedMember() {
+ if (nextMostLoadedMember < mostLoadedRangeStart) {
+ if (mostLoadedRangeEnd <= mostLoadedRangeStart &&
+ mostLoadedRangeStart > 0) {
+ // The range is empty due to calls to excludeMostLoadedMember(). We risk not
+ // expanding the range below and returning a member outside the range. Ensure
+ // that we always expand the range below by resetting the partition count.
+ mostLoadedRangePartitionCount = memberTargetAssignmentSizes[sortedMembers.get(mostLoadedRangeStart - 1)];
+ } else {
+ mostLoadedRangePartitionCount--;
}
+
+ // Expand the range.
+ while (mostLoadedRangeStart > 0 &&
+ memberTargetAssignmentSizes[sortedMembers.get(mostLoadedRangeStart - 1)] == mostLoadedRangePartitionCount) {
+ mostLoadedRangeStart--;
+ }
+
+ // Reset the pointer to the end of the range of members with identical partition
+ // counts.
+ nextMostLoadedMember = mostLoadedRangeEnd - 1;
}
- } 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 (groupSpec.memberSubscription(anotherMember).subscribedTopicIds().contains(partition.topicId())) {
- newOwner = anotherMember;
- break;
+ if (nextMostLoadedMember < 0) {
+ // The range is empty due to excludeMostLoadedMember() calls and there are no more
+ // members that can give up partitions.
+ return -1;
}
+
+ int mostLoadedMemberIndex = sortedMembers.get(nextMostLoadedMember);
+ nextMostLoadedMember--;
+
+ return mostLoadedMemberIndex;
}
- if (newOwner == null) {
- throw new PartitionAssignorException("No suitable new owner was found for the partition" + partition);
+ /**
+ * Excludes the last member returned from nextMostLoadedMember from the most loaded range.
+ *
+ * Must not be called if nextMostLoadedMember has not been called yet or returned -1.
+ */
+ public void excludeMostLoadedMember() {
+ // Kick the member out of the most loaded range by swapping with the member at the end
+ // of the range and contracting the end of the range.
+ int mostLoadedMemberIndex = sortedMembers.get(nextMostLoadedMember + 1);
+ sortedMembers.set(nextMostLoadedMember + 1, sortedMembers.get(mostLoadedRangeEnd - 1));
+ sortedMembers.set(mostLoadedRangeEnd - 1, mostLoadedMemberIndex);
+ mostLoadedRangeEnd--;
}
- reassignPartition(partition, newOwner);
+ /**
+ * Checks whether the members are balanced based on the partition counts of the least and
+ * most loaded ranges.
+ */
+ public boolean isBalanced() {
+ return mostLoadedRangePartitionCount - leastLoadedRangePartitionCount <= 1;
+ }
}
/**
- * Reassigns the given partition to a new member while considering partition movements and stickiness.
- * <p>
- * This method performs the following actions:
+ * Allocates the remaining unassigned partitions to members in a balanced manner.
* <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>
+ * <li>Topics are sorted to maximize the probability of a balanced assignment.</li>
+ * <li>Unassigned partitions within each topic are distributed to the least loaded
+ * subscribers, repeatedly.</li>
* </ol>
*
- * @param partition The {@link TopicIdPartition} to be reassigned.
- * @param newMember The Id of the member to which the partition should be reassigned.
+ * @param partitions The partitions to be assigned.
*/
- private void reassignPartition(TopicIdPartition partition, String newMember) {
- String member = partitionOwnerInTargetAssignment.get(partition);
- // Find the correct partition movement considering the stickiness requirement.
- TopicIdPartition partitionToBeMoved = partitionMovements.computeActualPartitionToBeMoved(
- partition,
- member,
- newMember
- );
- processPartitionMovement(partitionToBeMoved, newMember);
- }
+ private void assignRemainingPartitions(Map<Uuid, List<Integer>> partitions) {
+ List<Uuid> sortedTopicIds = sortTopicIds(partitions.keySet());
- private void processPartitionMovement(TopicIdPartition topicIdPartition, String newMember) {
- String oldMember = partitionOwnerInTargetAssignment.get(topicIdPartition);
+ MemberAssignmentBalancer memberAssignmentBalancer = new MemberAssignmentBalancer();
- partitionMovements.movePartition(topicIdPartition, oldMember, newMember);
+ for (Uuid topicId : sortedTopicIds) {
+ memberAssignmentBalancer.initialize(topicSubscribers.get(topicId));
- assignmentManager.removePartitionFromTargetAssignment(topicIdPartition, oldMember);
- assignmentManager.addPartitionToTargetAssignment(topicIdPartition, newMember);
+ for (int partition : partitions.get(topicId)) {
+ int leastLoadedMemberIndex = memberAssignmentBalancer.nextLeastLoadedMember();
+ assignPartition(topicId, partition, leastLoadedMemberIndex);
+ }
+ }
}
/**
- * Adds the topic's partition to the member's target assignment.
+ * 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 static void addPartitionToAssignment(
- Map<String, MemberAssignment> memberAssignments,
- String memberId,
- Uuid topicId,
- int partition
- ) {
- memberAssignments.get(memberId)
- .partitions()
- .computeIfAbsent(topicId, __ -> new HashSet<>())
- .add(partition);
+ private boolean canTopicParticipateInReassignment(Uuid topicId) {
+ return topicSubscribers.get(topicId).size() >= 2;
}
/**
- * Constructs a set of {@code TopicIdPartition} including all the given topic Ids based on their partition counts.
- *
- * @param topicIds Collection of topic Ids.
- * @param subscribedTopicDescriber Describer to fetch partition counts for topics.
- *
- * @return Set of {@code TopicIdPartition} including all the provided topic Ids.
+ * Balance the current assignment after the initial round of assignments have completed.
*/
- private static Set<TopicIdPartition> topicIdPartitions(
- Collection<Uuid> topicIds,
- SubscribedTopicDescriber subscribedTopicDescriber
- ) {
- Set<TopicIdPartition> topicIdPartitions = new HashSet<>();
- for (Uuid topicId : topicIds) {
- int numPartitions = subscribedTopicDescriber.numPartitions(topicId);
- for (int partitionId = 0; partitionId < numPartitions; partitionId++) {
- topicIdPartitions.add(new TopicIdPartition(topicId, partitionId));
+ private void balance() {
+ Set<Uuid> topicIds = new HashSet<>(subscribedTopicIds);
+
+ // Narrow down the reassignment scope to only those topics that can actually be reassigned.
+ for (Uuid topicId : subscribedTopicIds) {
+ if (!canTopicParticipateInReassignment(topicId)) {
+ topicIds.remove(topicId);
}
}
- return topicIdPartitions;
+
+ // If all the partitions are fixed i.e. unassigned partitions is empty there is no point of re-balancing.
+ if (!topicIds.isEmpty()) balanceTopics(topicIds);
}
/**
- * This class represents a pair of member Ids involved in a partition reassignment.
- * Each pair contains a source and a destination member Id.
- * It 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.
+ * Performs reassignments of partitions to balance the load across members.
+ * This method iteratively reassigns partitions until no further moves can improve the balance.
+ * <p/>
+ * The method loops over the topics repeatedly until an entire loop around the topics has
+ * completed without any reassignments, or we hit an iteration limit.
+ * <p/>
+ * This method produces perfectly balanced assignments when all subscribers of a topic have the
+ * same subscriptions. However, when subscribers have overlapping, non-identical subscriptions,
+ * the method produces almost-balanced assignments, assuming the iteration limit is not hit.
+ * eg. if there are three members and two topics and members 1 and 2 are subscribed to the first
+ * topic and members 2 and 3 are subscribed to the second topic, we can end up with an
+ * assignment like:
+ * <ul>
+ * <li>Member 1: 9 partitions</li>
+ * <li>Member 2: 10 partitions</li>
+ * <li>Member 3: 11 partitions</li>
+ * </ul>
+ *
+ * In this assignment, the subscribers of the first topic have a difference in partitions of 1,
+ * so the topic is considered balanced. The same applies to the second topic. However, balance
+ * can be improved by moving a partition from the second topic from member 3 to member 2 and a
+ * partition from the first topic from member 2 to member 1.
+ *
+ * @param topicIds The topics to consider for reassignment. These topics must have at
+ * least two subscribers.
*/
- 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;
+ private void balanceTopics(Collection<Uuid> topicIds) {
+ List<Uuid> sortedTopicIds = sortTopicIds(topicIds);
+ // The index of the last topic in sortedTopicIds that was rebalanced.
+ // Used to decide when to exit early.
+ int lastRebalanceTopicIndex = -1;
+
+ MemberAssignmentBalancer memberAssignmentBalancer = new MemberAssignmentBalancer();
+
+ // An array of partitions, with partitions owned by the same member grouped together.
+ List<Integer> partitions = new ArrayList<>();
+
+ // The ranges in the partitions list assigned to members.
+ // Maintaining these ranges allows for constant time, deterministic picking of partitions
+ // owned by a given member.
+ Map<Integer, Integer> startPartitionIndices = new HashMap<>();
+ Map<Integer, Integer> endPartitionIndices = new HashMap<>(); // exclusive
+
+ // Repeat reassignment until no partition can be moved to improve the balance or we hit an
+ // iteration limit.
+ for (int i = 0; i < MAX_ITERATION_COUNT; i++) {
+ for (int topicIndex = 0; topicIndex < sortedTopicIds.size(); topicIndex++) {
+ if (topicIndex == lastRebalanceTopicIndex) {
+ // The last rebalanced topic was this one, which means we've gone through all
+ // the topics and didn't perform any additional rebalancing. Don't bother trying
+ // to rebalance this topic again and exit.
+ return;
+ }
- if (!getClass().isInstance(obj))
- return false;
+ Uuid topicId = sortedTopicIds.get(topicIndex);
- MemberPair otherPair = (MemberPair) obj;
- return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
- }
+ int reassignedPartitionCount = balanceTopic(
+ topicId,
+ memberAssignmentBalancer,
+ partitions,
+ startPartitionIndices,
+ endPartitionIndices
+ );
- @Override
- public String toString() {
- return "MemberPair(" +
- "srcMemberId='" + srcMemberId + '\'' +
- ", dstMemberId='" + dstMemberId + '\'' +
- ')';
+ if (reassignedPartitionCount > 0 ||
+ lastRebalanceTopicIndex == -1) {
+ lastRebalanceTopicIndex = topicIndex;
+ }
+ }
}
}
/**
- * 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.
+ * Performs reassignments of a topic's partitions to balance the load across its subscribers.
+ *
+ * @param topicId The topic to consider for reassignment. The topic must have
+ * at least two subscribers.
+ * @param memberAssignmentBalancer A MemberAssignmentBalancer to hold the state of the balancing
+ * algorithm.
+ * subscribers in ascending order of total assigned partitions.
+ * @param partitions A list for the balancing algorithm to store the topic's
+ * partitions, with partitions owned by the same member grouped
+ * together.
+ * @param startPartitionIndices A map for the balancing algorithm to store the ranges in the
+ * partitions list assigned to members. These ranges allows for
+ * constant time, deterministic picking of partitions owned by a
+ * given member.
+ * @param endPartitionIndices A map for the balancing algorithm to store the ranges in the
+ * partitions list assigned to members. These ranges allows for
+ * constant time, deterministic picking of partitions owned by a
+ * given member.
+ * @return the number of partitions reassigned.
*/
- 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 int balanceTopic(
+ Uuid topicId,
+ MemberAssignmentBalancer memberAssignmentBalancer,
+ List<Integer> partitions,
+ Map<Integer, Integer> startPartitionIndices,
+ Map<Integer, Integer> endPartitionIndices // exclusive
+ ) {
+ int reassignedPartitionCount = 0;
- private void addPartitionMovementRecord(TopicIdPartition partition, MemberPair pair) {
- partitionMovementsByPartition.put(partition, pair);
+ int numPartitions = subscribedTopicDescriber.numPartitions(topicId);
- Uuid topic = partition.topicId();
- if (!partitionMovementsByTopic.containsKey(topic))
- partitionMovementsByTopic.put(topic, new HashMap<>());
+ int[] partitionOwners = targetAssignmentPartitionOwners.get(topicId);
- Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
- if (!partitionMovementsForThisTopic.containsKey(pair))
- partitionMovementsForThisTopic.put(pair, new HashSet<>());
+ // Next, try to reassign partitions from the most loaded subscribers to the least loaded
+ // subscribers. We use a similar approach to assignRemainingPartitions, except instead
+ // we are reassigning partitions from the most loaded subscribers to the least loaded
+ // subscribers. Once the difference in partition counts between the ranges is one or less,
+ // we are done, since reassigning partitions would not improve balance.
+ int imbalance = memberAssignmentBalancer.initialize(topicSubscribers.get(topicId));
+ if (imbalance <= 1) {
+ // The topic is already balanced.
+ return reassignedPartitionCount;
+ }
- partitionMovementsForThisTopic.get(pair).add(partition);
+ // Initialize the array of partitions, with partitions owned by the same member grouped
+ // together.
+ partitions.clear();
+ for (int partition = 0; partition < numPartitions; partition++) {
+ partitions.add(partition);
}
+ partitions.sort(
+ Comparator
+ .comparingInt((Integer partition) -> partitionOwners[partition])
+ .thenComparingInt(partition -> partition)
+ );
- private void movePartition(TopicIdPartition partition, String oldOwner, String newOwner) {
- MemberPair pair = new MemberPair(oldOwner, newOwner);
+ // Initialize the ranges in the partitions list owned by members.
+ startPartitionIndices.clear();
+ endPartitionIndices.clear();
- if (partitionMovementsByPartition.containsKey(partition)) {
- // This partition was previously moved.
- MemberPair existingPair = removeMovementRecordOfPartition(partition);
- if (existingPair.dstMemberId.equals(oldOwner)) {
- throw new PartitionAssignorException("Mismatch in partition movement record with respect to " +
- "partition ownership during a rebalance"
- );
- }
- if (!existingPair.srcMemberId.equals(newOwner)) {
- // The partition is not moving back to its previous member.
- addPartitionMovementRecord(partition, new MemberPair(existingPair.srcMemberId, newOwner));
- }
- } else
- addPartitionMovementRecord(partition, pair);
+ for (int i = 0; i < numPartitions; i++) {
+ int partition = partitions.get(i);
+ int ownerIndex = partitionOwners[partition];
+ startPartitionIndices.putIfAbsent(ownerIndex, i);
+ endPartitionIndices.put(ownerIndex, i + 1); // endPartitionIndices is exclusive
}
- /**
- * Computes the actual partition to be moved based on the current and proposed partition owners.
- * This method determines the appropriate partition movement, considering existing partition movements
- * and constraints within a topic.
- *
- * @param partition The {@link TopicIdPartition} object representing the partition to be moved.
- * @param oldOwner The memberId of the current owner of the partition.
- * @param newOwner The memberId of the proposed new owner of the partition.
- * @return The {@link TopicIdPartition} that should be moved, based on existing movement patterns
- * and ownership. Returns the original partition if no specific movement pattern applies.
- * @throws PartitionAssignorException if the old owner does not match the expected value for the partition.
- */
- private TopicIdPartition computeActualPartitionToBeMoved(
- 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;
- }
+ // Redistribute partitions from the most loaded subscriber to the least loaded subscriber
+ // until the topic's subscribers are balanced.
+ while (!memberAssignmentBalancer.isBalanced()) {
+ // NB: The condition above does not do much. This loop will terminate due to the checks
+ // inside instead.
- Map<MemberPair, Set<TopicIdPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
- MemberPair reversePair = new MemberPair(newOwner, oldOwner);
- if (!partitionMovementsForThisTopic.containsKey(reversePair))
- return partition;
+ // First, choose a member from the most loaded range to reassign a partition from.
- return partitionMovementsForThisTopic.get(reversePair).iterator().next();
- }
- }
+ // Loop until we find a member that has partitions to give up.
+ int mostLoadedMemberIndex = -1;
+ while (true) {
+ mostLoadedMemberIndex = memberAssignmentBalancer.nextMostLoadedMember();
- /**
- * 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<>();
+ if (mostLoadedMemberIndex == -1) {
+ // There are no more members with partitions to give up.
+ // We need to break out of two loops here.
+ break;
+ }
- /**
- * 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;
- }
+ if (!endPartitionIndices.containsKey(mostLoadedMemberIndex) ||
+ endPartitionIndices.get(mostLoadedMemberIndex) - startPartitionIndices.get(mostLoadedMemberIndex) <= 0) {
+ memberAssignmentBalancer.excludeMostLoadedMember();
+ continue;
+ }
- @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);
+ break;
}
- @Override
- public int hashCode() {
- return Objects.hash(memberId);
+ if (mostLoadedMemberIndex == -1) {
+ // There are no more members with partitions to give up.
+ break;
}
- @Override
- public String toString() {
- return "MemberAssignmentData(" +
- "memberId='" + memberId + '\'' +
- ", currentAssignmentSize=" + currentAssignmentSize +
- ", maxAssignmentSize=" + maxAssignmentSize +
- ')';
- }
- }
+ // We've picked a member. Now pick its last partition to reassign.
+ int partition = partitions.get(endPartitionIndices.get(mostLoadedMemberIndex) - 1);
+ endPartitionIndices.put(mostLoadedMemberIndex, endPartitionIndices.get(mostLoadedMemberIndex) - 1);
- /**
- * Initializes an AssignmentManager, setting up the necessary data structures.
- */
- public AssignmentManager(
- SubscribedTopicDescriber subscribedTopicDescriber
- ) {
- groupSpec.memberIds().forEach(memberId -> {
- int maxSize = groupSpec.memberSubscription(memberId).subscribedTopicIds().stream()
- .mapToInt(subscribedTopicDescriber::numPartitions)
- .sum();
-
- MemberAssignmentData memberAssignmentData = membersWithAssignmentSizes
- .computeIfAbsent(memberId, MemberAssignmentData::new);
- memberAssignmentData.maxAssignmentSize = maxSize;
- memberAssignmentData.currentAssignmentSize = 0;
- });
- }
+ // Find the least loaded subscriber.
+ int leastLoadedMemberIndex = memberAssignmentBalancer.nextLeastLoadedMember();
- /**
- * @param memberId The member Id.
- * @return The current assignment size for the given member.
- */
- private int targetAssignmentSize(String memberId) {
- MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
- if (memberData == null) {
- LOG.warn("Member Id {} not found", memberId);
- return 0;
+ // If the most and least loaded subscribers came from member ranges with partition
+ // counts that differ by one or less, then the members are balanced and we are done.
+ if (memberAssignmentBalancer.isBalanced()) {
+ // The topic is balanced.
+ break;
}
- return memberData.currentAssignmentSize;
- }
- /**
- * @param memberId The member Id.
- * @return The maximum assignment size for the given member.
- */
- private 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;
+ // Reassign the partition.
+ assignPartition(topicId, partition, leastLoadedMemberIndex);
+ reassignedPartitionCount++;
}
- /**
- * @param memberId The member Id.
- * @return If the given member is at maximum capacity.
- */
- private boolean isMemberAtMaxCapacity(String memberId) {
- return targetAssignmentSize(memberId) >= maxAssignmentSize(memberId);
- }
+ return reassignedPartitionCount;
+ }
- /**
- * @param memberId The member Id.
- * Increment the current target assignment size for the member.
- */
- private void incrementTargetAssignmentSize(String memberId) {
- MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
- if (memberData == null) {
- LOG.warn("Member Id {} not found", memberId);
- return;
- }
- memberData.currentAssignmentSize++;
- }
+ /**
+ * Computes the set of unassigned partitions, based on targetAssignmentPartitionOwners.
+ */
+ private Map<Uuid, List<Integer>> computeUnassignedPartitions() {
+ Map<Uuid, List<Integer>> unassignedPartitions = new HashMap<>();
- /**
- * @param memberId The member Id.
- * Decrement the current target assignment size for the member, if it's assignment size is greater than zero.
- */
- private void decrementTargetAssignmentSize(String memberId) {
- MemberAssignmentData memberData = this.membersWithAssignmentSizes.get(memberId);
- if (memberData == null) {
- LOG.warn("Member Id {} not found", memberId);
- return;
+ for (Uuid topicId : subscribedTopicIds) {
+ List<Integer> topicUnassignedPartitions = new ArrayList<>();
+ int numPartitions = subscribedTopicDescriber.numPartitions(topicId);
+ for (int partition = 0; partition < numPartitions; partition++) {
+ if (targetAssignmentPartitionOwners.get(topicId)[partition] == -1) {
+ topicUnassignedPartitions.add(partition);
+ }
}
- if (memberData.currentAssignmentSize > 0) {
- memberData.currentAssignmentSize--;
+ if (!topicUnassignedPartitions.isEmpty()) {
+ unassignedPartitions.put(topicId, topicUnassignedPartitions);
}
}
- /**
- * 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.
- */
- private boolean maybeAssignPartitionToMember(
- TopicIdPartition topicIdPartition,
- String memberId
- ) {
- // If member is not subscribed to the partition's topic, return false without assigning.
- if (!groupSpec.memberSubscription(memberId).subscribedTopicIds().contains(topicIdPartition.topicId())) {
- return false;
- }
+ return unassignedPartitions;
+ }
- // If the member's current assignment is already at max, return false without assigning.
- if (isMemberAtMaxCapacity(memberId)) {
- return false;
- }
+ /**
+ * Assigns or reassigns the given partition to a member.
+ *
+ * @param topicId The topic containing the partition to be assigned or reassigned.
+ * @param partition The partition to be assigned or reassigned.
+ * @param memberIndex The index of the member to which the partition should be assigned.
+ */
+ private void assignPartition(Uuid topicId, int partition, int memberIndex) {
+ int oldMemberIndex = targetAssignmentPartitionOwners.get(topicId)[partition];
- addPartitionToTargetAssignment(topicIdPartition, memberId);
- return true;
+ if (oldMemberIndex != -1) {
+ removePartitionFromTargetAssignment(topicId, partition, oldMemberIndex);
}
- /**
- * Assigns a partition to a member, 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 addPartitionToTargetAssignment(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 (!isMemberAtMaxCapacity(memberId)) {
- sortedMembersByAssignmentSize.add(memberId);
- }
+ addPartitionToTargetAssignment(topicId, partition, memberIndex);
+ }
+
+ /**
+ * @param memberIndex The member index.
+ * @return The current assignment size for the given member.
+ */
+ private int targetAssignmentSize(int memberIndex) {
+ return memberTargetAssignmentSizes[memberIndex];
+ }
- unassignedPartitions.remove(topicIdPartition);
+ /**
+ * Assigns a partition to a member and updates the current assignment size.
+ *
+ * @param topicId The topic containing the partition to be assigned.
+ * @param partition The partition to be assigned.
+ * @param memberIndex Member that the partition needs to be added to.
+ */
+ private void addPartitionToTargetAssignment(
+ Uuid topicId,
+ int partition,
+ int memberIndex
+ ) {
+ String memberId = memberIds.get(memberIndex);
+ Map<Uuid, Set<Integer>> assignment = targetAssignment.get(memberId).partitions();
+ if (AssignorHelpers.isImmutableMap(assignment)) {
+ assignment = AssignorHelpers.deepCopyAssignment(assignment);
+ targetAssignment.put(memberId, new MemberAssignmentImpl(assignment));
}
+ assignment
+ .computeIfAbsent(topicId, __ -> {
+ int numPartitions = subscribedTopicDescriber.numPartitions(topicId);
+ int numSubscribers = topicSubscribers.get(topicId).size();
+ int estimatedPartitionsPerSubscriber = (numPartitions + numSubscribers - 1) / numSubscribers;
+ return new HashSet<>((int) ((estimatedPartitionsPerSubscriber / 0.75f) + 1));
+ })
+ .add(partition);
- /**
- * 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 removePartitionFromTargetAssignment(TopicIdPartition topicIdPartition, String memberId) {
- Map<Uuid, Set<Integer>> targetPartitionsMap = targetAssignment.get(memberId).partitions();
- 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());
- }
- }
+ targetAssignmentPartitionOwners.get(topicId)[partition] = memberIndex;
- partitionOwnerInTargetAssignment.remove(topicIdPartition, memberId);
- // Remove the member's assignment data from the set to update it.
- sortedMembersByAssignmentSize.remove(memberId);
- assignmentManager.decrementTargetAssignmentSize(memberId);
+ memberTargetAssignmentSizes[memberIndex]++;
+ }
- // Update current assignment size and re-add to set if needed.
- if (!isMemberAtMaxCapacity(memberId)) {
- sortedMembersByAssignmentSize.add(memberId);
+ /**
+ * Revokes the partition from a member and updates the current target assignment size.
+ *
+ * @param topicId The topic containing the partition to be revoked.
+ * @param partition The partition to be revoked.
+ * @param memberIndex Member that the partition needs to be revoked from.
+ */
+ private void removePartitionFromTargetAssignment(
+ Uuid topicId,
+ int partition,
+ int memberIndex
+ ) {
+ String memberId = memberIds.get(memberIndex);
+ Map<Uuid, Set<Integer>> assignment = targetAssignment.get(memberId).partitions();
+ if (AssignorHelpers.isImmutableMap(assignment)) {
+ assignment = AssignorHelpers.deepCopyAssignment(assignment);
+ targetAssignment.put(memberId, new MemberAssignmentImpl(assignment));
+ }
+ Set<Integer> partitionsSet = assignment.get(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(partition);
+ if (partitionsSet.isEmpty()) {
+ assignment.remove(topicId);
}
}
- /**
- * Sorts members in ascending order based on their current target assignment size.
- * Members that have reached their max assignment size are removed.
- *
- * @param memberIds Member Ids that need to be sorted.
- * @return A set that maintains the order of members by assignment size.
- */
- private TreeSet<String> sortMembersByAssignmentSize(Collection<String> memberIds) {
- Comparator<String> comparator = Comparator
- .comparingInt((String memberId) -> membersWithAssignmentSizes.get(memberId).currentAssignmentSize)
- .thenComparing(memberId -> memberId);
-
- return memberIds.stream()
- .filter(memberId -> {
- MemberAssignmentData memberData = membersWithAssignmentSizes.get(memberId);
- return memberData.currentAssignmentSize < memberData.maxAssignmentSize;
- })
- .collect(Collectors.toCollection(() -> new TreeSet<>(comparator)));
- }
+ targetAssignmentPartitionOwners.get(topicId)[partition] = -1;
+
+ memberTargetAssignmentSizes[memberIndex]--;
}
}
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformHomogeneousAssignmentBuilder.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformHomogeneousAssignmentBuilder.java
index b4b8a306efc..8e7fe209c82 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformHomogeneousAssignmentBuilder.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/assignor/UniformHomogeneousAssignmentBuilder.java
@@ -49,17 +49,6 @@ import java.util.Set;
* Balance > Stickiness.
*/
public class UniformHomogeneousAssignmentBuilder {
- private static final Class<?> UNMODIFIABLE_MAP_CLASS = Collections.unmodifiableMap(new HashMap<>()).getClass();
- private static final Class<?> EMPTY_MAP_CLASS = Collections.emptyMap().getClass();
-
- /**
- * @return True if the provided map is an UnmodifiableMap or EmptyMap. Those classes are not
- * public hence we cannot use the `instanceof` operator.
- */
- private static boolean isImmutableMap(Map<?, ?> map) {
- return UNMODIFIABLE_MAP_CLASS.isInstance(map) || EMPTY_MAP_CLASS.isInstance(map);
- }
-
/**
* The assignment specification which includes member metadata.
*/
@@ -171,7 +160,7 @@ public class UniformHomogeneousAssignmentBuilder {
// The assignor expects to receive the assignment as an immutable map. It leverages
// this knowledge in order to avoid having to copy all assignments.
- if (!isImmutableMap(oldAssignment)) {
+ if (!AssignorHelpers.isImmutableMap(oldAssignment)) {
throw new IllegalStateException("The assignor expect an immutable map.");
}
@@ -196,7 +185,7 @@ public class UniformHomogeneousAssignmentBuilder {
if (newAssignment == null) {
// If the new assignment is null, we create a deep copy of the
// original assignment so that we can alter it.
- newAssignment = deepCopy(oldAssignment);
+ newAssignment = AssignorHelpers.deepCopyAssignment(oldAssignment);
}
// Remove the partition from the new assignment.
Set<Integer> parts = newAssignment.get(topicId);
@@ -213,7 +202,7 @@ public class UniformHomogeneousAssignmentBuilder {
if (newAssignment == null) {
// If the new assignment is null, we create a deep copy of the
// original assignment so that we can alter it.
- newAssignment = deepCopy(oldAssignment);
+ newAssignment = AssignorHelpers.deepCopyAssignment(oldAssignment);
}
// Remove the entire topic.
newAssignment.remove(topicId);
@@ -243,10 +232,10 @@ public class UniformHomogeneousAssignmentBuilder {
int remainingQuota = unfilledMember.remainingQuota;
Map<Uuid, Set<Integer>> newAssignment = targetAssignment.get(memberId).partitions();
- if (isImmutableMap(newAssignment)) {
+ if (AssignorHelpers.isImmutableMap(newAssignment)) {
// If the new assignment is immutable, we must create a deep copy of it
// before altering it.
- newAssignment = deepCopy(newAssignment);
+ newAssignment = AssignorHelpers.deepCopyAssignment(newAssignment);
targetAssignment.put(memberId, new MemberAssignmentImpl(newAssignment));
}
@@ -264,14 +253,6 @@ public class UniformHomogeneousAssignmentBuilder {
}
}
- private static Map<Uuid, Set<Integer>> deepCopy(Map<Uuid, Set<Integer>> map) {
- Map<Uuid, Set<Integer>> copy = new HashMap<>(map.size());
- for (Map.Entry<Uuid, Set<Integer>> entry : map.entrySet()) {
- copy.put(entry.getKey(), new HashSet<>(entry.getValue()));
- }
- return copy;
- }
-
private static class MemberWithRemainingQuota {
final String memberId;
final int remainingQuota;
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java
index 521414b4392..22425476c6f 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java
@@ -20,6 +20,7 @@ import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignorException;
+import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType;
import org.apache.kafka.coordinator.group.modern.Assignment;
import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
@@ -28,8 +29,11 @@ import org.apache.kafka.coordinator.group.modern.TopicMetadata;
import org.junit.jupiter.api.Test;
+import java.util.ArrayList;
+import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
@@ -57,6 +61,33 @@ public class UniformHeterogeneousAssignmentBuilderTest {
private final String memberB = "B";
private final String memberC = "C";
+ /**
+ * A GroupSpec implementation that returns members in sorted order, so that assignor output is
+ * deterministic.
+ */
+ private static class TestGroupSpecImpl extends GroupSpecImpl {
+ private final List<String> orderedMembers;
+
+ public TestGroupSpecImpl(
+ Map<String, MemberSubscriptionAndAssignmentImpl> members,
+ SubscriptionType subscriptionType,
+ Map<Uuid, Map<Integer, String>> invertedMemberAssignment
+ ) {
+ super(members, subscriptionType, invertedMemberAssignment);
+
+ this.orderedMembers = new ArrayList<>(members.keySet());
+ this.orderedMembers.sort(null);
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public Collection<String> memberIds() {
+ return orderedMembers;
+ }
+ }
+
@Test
public void testTwoMembersNoTopicSubscription() {
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
@@ -84,7 +115,7 @@ public class UniformHeterogeneousAssignmentBuilderTest {
Assignment.EMPTY
));
- GroupSpec groupSpec = new GroupSpecImpl(
+ GroupSpec groupSpec = new TestGroupSpecImpl(
members,
HETEROGENEOUS,
Collections.emptyMap()
@@ -125,7 +156,7 @@ public class UniformHeterogeneousAssignmentBuilderTest {
Assignment.EMPTY
));
- GroupSpec groupSpec = new GroupSpecImpl(
+ GroupSpec groupSpec = new TestGroupSpecImpl(
members,
HETEROGENEOUS,
Collections.emptyMap()
@@ -166,7 +197,7 @@ public class UniformHeterogeneousAssignmentBuilderTest {
Assignment.EMPTY
));
- GroupSpec groupSpec = new GroupSpecImpl(
+ GroupSpec groupSpec = new TestGroupSpecImpl(
members,
HETEROGENEOUS,
Collections.emptyMap()
@@ -181,10 +212,10 @@ public class UniformHeterogeneousAssignmentBuilderTest {
Map<String, Map<Uuid, Set<Integer>>> expectedAssignment = new HashMap<>();
expectedAssignment.put(memberA, mkAssignment(
mkTopicAssignment(topic1Uuid, 0, 1, 2),
- mkTopicAssignment(topic3Uuid, 3, 5)
+ mkTopicAssignment(topic3Uuid, 4)
));
expectedAssignment.put(memberB, mkAssignment(
- mkTopicAssignment(topic3Uuid, 0, 1, 2, 4)
+ mkTopicAssignment(topic3Uuid, 0, 1, 2, 3, 5)
));
assertAssignment(expectedAssignment, computedAssignment);
@@ -227,7 +258,7 @@ public class UniformHeterogeneousAssignmentBuilderTest {
Assignment.EMPTY
));
- GroupSpec groupSpec = new GroupSpecImpl(
+ GroupSpec groupSpec = new TestGroupSpecImpl(
members,
HETEROGENEOUS,
Collections.emptyMap()
@@ -305,7 +336,7 @@ public class UniformHeterogeneousAssignmentBuilderTest {
))
));
- GroupSpec groupSpec = new GroupSpecImpl(
+ GroupSpec groupSpec = new TestGroupSpecImpl(
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
@@ -319,13 +350,14 @@ public class UniformHeterogeneousAssignmentBuilderTest {
Map<String, Map<Uuid, Set<Integer>>> expectedAssignment = new HashMap<>();
expectedAssignment.put(memberA, mkAssignment(
- mkTopicAssignment(topic1Uuid, 0, 1, 2, 3, 4)
+ mkTopicAssignment(topic1Uuid, 0, 1, 2, 5)
));
expectedAssignment.put(memberB, mkAssignment(
+ mkTopicAssignment(topic1Uuid, 3),
mkTopicAssignment(topic2Uuid, 0, 1, 2, 3)
));
expectedAssignment.put(memberC, mkAssignment(
- mkTopicAssignment(topic1Uuid, 5),
+ mkTopicAssignment(topic1Uuid, 4),
mkTopicAssignment(topic3Uuid, 0, 1, 2, 3)
));
@@ -379,7 +411,7 @@ public class UniformHeterogeneousAssignmentBuilderTest {
))
));
- GroupSpec groupSpec = new GroupSpecImpl(
+ GroupSpec groupSpec = new TestGroupSpecImpl(
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
@@ -449,7 +481,7 @@ public class UniformHeterogeneousAssignmentBuilderTest {
Assignment.EMPTY
));
- GroupSpec groupSpec = new GroupSpecImpl(
+ GroupSpec groupSpec = new TestGroupSpecImpl(
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
@@ -467,10 +499,10 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
expectedAssignment.put(memberB, mkAssignment(
mkTopicAssignment(topic1Uuid, 1),
- mkTopicAssignment(topic2Uuid, 1, 2, 5)
+ mkTopicAssignment(topic2Uuid, 1, 2, 6)
));
expectedAssignment.put(memberC, mkAssignment(
- mkTopicAssignment(topic2Uuid, 0, 3, 4, 6)
+ mkTopicAssignment(topic2Uuid, 0, 3, 4, 5)
));
assertAssignment(expectedAssignment, computedAssignment);
@@ -518,7 +550,7 @@ public class UniformHeterogeneousAssignmentBuilderTest {
// Member C was removed
- GroupSpec groupSpec = new GroupSpecImpl(
+ GroupSpec groupSpec = new TestGroupSpecImpl(
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
@@ -579,7 +611,7 @@ public class UniformHeterogeneousAssignmentBuilderTest {
))
));
- GroupSpec groupSpec = new GroupSpecImpl(
+ GroupSpec groupSpec = new TestGroupSpecImpl(
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
@@ -602,6 +634,76 @@ public class UniformHeterogeneousAssignmentBuilderTest {
assertAssignment(expectedAssignment, computedAssignment);
}
+ /**
+ * Tests an edge case in the reassignment algorithm where we run out of topic partitions before
+ * the topic's subscribers are balanced, and the topic's least loaded subscribers start off with
+ * no partitions from the topic. The per-topic reassignment loop should break when it cannot
+ * find a partition to reassign.
+ */
+ @Test
+ public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions() {
+ Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
+ topicMetadata.put(topic1Uuid, new TopicMetadata(
+ topic1Uuid,
+ topic1Name,
+ 2
+ ));
+ topicMetadata.put(topic2Uuid, new TopicMetadata(
+ topic2Uuid,
+ topic2Name,
+ 2
+ ));
+ topicMetadata.put(topic3Uuid, new TopicMetadata(
+ topic3Uuid,
+ topic3Name,
+ 2
+ ));
+
+ Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
+
+ members.put(memberA, new MemberSubscriptionAndAssignmentImpl(
+ Optional.empty(),
+ Optional.empty(),
+ Set.of(topic1Uuid, topic2Uuid, topic3Uuid),
+ new Assignment(mkAssignment(
+ mkTopicAssignment(topic1Uuid, 0, 1),
+ mkTopicAssignment(topic2Uuid, 0, 1),
+ mkTopicAssignment(topic3Uuid, 0, 1)
+ ))
+ ));
+
+ members.put(memberB, new MemberSubscriptionAndAssignmentImpl(
+ Optional.empty(),
+ Optional.empty(),
+ Set.of(topic1Uuid, topic2Uuid, topic3Uuid),
+ Assignment.EMPTY
+ ));
+
+ GroupSpec groupSpec = new TestGroupSpecImpl(
+ members,
+ HETEROGENEOUS,
+ invertedTargetAssignment(members)
+ );
+ SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+
+ GroupAssignment computedAssignment = assignor.assign(
+ groupSpec,
+ subscribedTopicMetadata
+ );
+
+ Map<String, Map<Uuid, Set<Integer>>> expectedAssignment = new HashMap<>();
+ expectedAssignment.put(memberA, mkAssignment(
+ mkTopicAssignment(topic2Uuid, 0),
+ mkTopicAssignment(topic3Uuid, 0, 1)
+ ));
+ expectedAssignment.put(memberB, mkAssignment(
+ mkTopicAssignment(topic1Uuid, 0, 1),
+ mkTopicAssignment(topic2Uuid, 1)
+ ));
+
+ assertAssignment(expectedAssignment, computedAssignment);
+ }
+
@Test
public void testFirstAssignmentWithTwoMembersIncludingOneWithoutSubscriptions() {
Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
@@ -627,7 +729,7 @@ public class UniformHeterogeneousAssignmentBuilderTest {
Assignment.EMPTY
));
- GroupSpec groupSpec = new GroupSpecImpl(
+ GroupSpec groupSpec = new TestGroupSpecImpl(
members,
HETEROGENEOUS,
Collections.emptyMap()