KAFKA-21318, KAFKA-20292: Compute GroupSpecImpl's assigned partitions lazily (#23883)

Replace GroupSpecImpl's inverted target assignment with a BitSet of the
assigned partitions of each topic, computed from the members' current
assignments on the first call to isPartitionAssigned. Only the uniform
homogeneous assignor calls isPartitionAssigned, so other assignors no
longer compute it.

This makes build in TargetAssignmentBuilderBenchmark ~40% faster for
homogeneous groups, because isPartitionAssigned no longer looks up
partitions in the inverted target assignment's TimelineHashMaps, which
are slower than regular maps.

It also fixes a bug where the assignor saw the partitions of members
removed in the same request as still assigned. The bug is not reachable
today, since we do not compute a new target assignment when a member
leaves, only on the next heartbeat, and static member replacement keeps
the replaced member's partitions assigned.

Remove ModernGroup's inverted target assignment, since it is no longer
read. For share groups, it was never correct, because a partition can be
assigned to several members of a share group, while the inverted target
assignment holds a single member per partition.

Rename buildAssignment in TargetAssignmentBuilderBenchmark to
buildGroupSpecAndAssignment, which builds a new group spec on every
invocation, so that the cost of computing the assigned partitions is
included, and replace buildGroupSpecWithoutOffload and
buildGroupSpecWithOffload with an assignorOffload parameter.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Reviewers: David Jacot <dajac@apache.org>
This commit is contained in:
Sean QuahandClaude Opus 5.5 authored and GitHub committed 2026-10-23 09:44:17 +01:00
1 parent da22240966
commit b860fd33a7
20 files changed
+383 -751

No files matched your search

@@ -4384,7 +4384,6 @@ public class GroupMetadataManager {
.withMembers(updatedMembersAndTargetAssignment.members())
.withSubscriptionType(subscriptionType)
.withTargetAssignment(updatedMembersAndTargetAssignment.targetAssignment())
.withInvertedTargetAssignment(group.invertedTargetAssignment())
.withMetadataImage(metadataImage)
.withResolvedRegularExpressions(group.resolvedRegularExpressions())
.build();
@@ -4480,7 +4479,6 @@ public class GroupMetadataManager {
.withSubscriptionType(subscriptionType)
.withTargetAssignment(updatedMembersAndTargetAssignment.targetAssignment())
.withTopicAssignablePartitionsMap(initializedTopicPartitions)
.withInvertedTargetAssignment(group.invertedTargetAssignment())
.withMetadataImage(metadataImage)
.build();
@@ -131,12 +131,6 @@ public abstract class GroupSpecBuilder<T extends ModernGroupMember, U extends Gr
*/
private Map<String, Assignment> targetAssignment = Map.of();
/**
* Reverse lookup map representing topic partitions with
* their current member assignments.
*/
private Map<Uuid, Map<Integer, String>> invertedTargetAssignment = Map.of();
/**
* The metadata image.
*/
@@ -193,19 +187,6 @@ public abstract class GroupSpecBuilder<T extends ModernGroupMember, U extends Gr
return self();
}
/**
* Adds the existing topic partition assignments.
*
* @param invertedTargetAssignment The reverse lookup map of the current target assignment.
* @return This object.
*/
public U withInvertedTargetAssignment(
Map<Uuid, Map<Integer, String>> invertedTargetAssignment
) {
this.invertedTargetAssignment = invertedTargetAssignment;
return self();
}
/**
* Adds the metadata image.
*
@@ -261,15 +242,8 @@ public abstract class GroupSpecBuilder<T extends ModernGroupMember, U extends Gr
))
);
Map<Uuid, Map<Integer, String>> invertedTargetAssignment = this.invertedTargetAssignment;
Optional<Map<Uuid, Set<Integer>>> topicAssignablePartitionsMap = this.topicAssignablePartitionsMap;
if (assignorOffload) {
Map<Uuid, Map<Integer, String>> invertedTargetAssignmentCopy = Utils.newHashMap(invertedTargetAssignment.size());
for (Map.Entry<Uuid, Map<Integer, String>> entry : invertedTargetAssignment.entrySet()) {
invertedTargetAssignmentCopy.put(entry.getKey(), Map.copyOf(entry.getValue()));
}
invertedTargetAssignment = invertedTargetAssignmentCopy;
if (topicAssignablePartitionsMap.isPresent()) {
Map<Uuid, Set<Integer>> topicAssignablePartitionsMapCopy = Utils.newHashMap(topicAssignablePartitionsMap.get().size());
for (Map.Entry<Uuid, Set<Integer>> entry : topicAssignablePartitionsMap.get().entrySet()) {
@@ -280,9 +254,9 @@ public abstract class GroupSpecBuilder<T extends ModernGroupMember, U extends Gr
}
return new GroupSpecImpl(
metadataImage,
Collections.unmodifiableMap(memberSpecs),
subscriptionType,
invertedTargetAssignment,
topicAssignablePartitionsMap
);
}
@@ -17,12 +17,15 @@
package org.apache.kafka.coordinator.group.modern;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
import org.apache.kafka.coordinator.group.api.assignor.MemberAssignment;
import org.apache.kafka.coordinator.group.api.assignor.MemberSubscription;
import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType;
import java.util.BitSet;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
@@ -32,6 +35,19 @@ import java.util.Set;
* The assignment specification for a modern group.
*/
public class GroupSpecImpl implements GroupSpec {
/**
* The partitions of a topic which are currently assigned to a member.
*
* @param partitions The assigned partitions.
* @param partitionCount The number of partitions in the topic.
*/
private record AssignedPartitions(BitSet partitions, int partitionCount) { }
/**
* The metadata image, used to size the assigned partitions of each topic.
*/
private final CoordinatorMetadataImage metadataImage;
/**
* Member subscription metadata keyed by member Id.
*/
@@ -43,10 +59,10 @@ public class GroupSpecImpl implements GroupSpec {
private final SubscriptionType subscriptionType;
/**
* Reverse lookup map representing topic partitions with
* their current member assignments.
* The partitions which are currently assigned to a member,
* keyed by topic id.
*/
private final Map<Uuid, Map<Integer, String>> invertedMemberAssignment;
private Map<Uuid, AssignedPartitions> assignedTopicPartitions;
/**
* In case of share groups, this map will be queried to decide
@@ -56,23 +72,23 @@ public class GroupSpecImpl implements GroupSpec {
private final Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap;
public GroupSpecImpl(
CoordinatorMetadataImage metadataImage,
Map<String, MemberSubscriptionAndAssignmentImpl> members,
SubscriptionType subscriptionType,
Map<Uuid, Map<Integer, String>> invertedMemberAssignment
SubscriptionType subscriptionType
) {
this(members, subscriptionType, invertedMemberAssignment, Optional.empty());
this(metadataImage, members, subscriptionType, Optional.empty());
}
public GroupSpecImpl(
CoordinatorMetadataImage metadataImage,
Map<String, MemberSubscriptionAndAssignmentImpl> members,
SubscriptionType subscriptionType,
Map<Uuid, Map<Integer, String>> invertedMemberAssignment,
Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap
) {
this.members = Objects.requireNonNull(members);
this.subscriptionType = Objects.requireNonNull(subscriptionType);
this.invertedMemberAssignment = Objects.requireNonNull(invertedMemberAssignment);
this.topicPartitionAllowedMap = Objects.requireNonNull(topicPartitionAllowedMap);
this.metadataImage = Objects.requireNonNull(metadataImage);
}
/**
@@ -93,14 +109,43 @@ public class GroupSpecImpl implements GroupSpec {
/**
* {@inheritDoc}
* <p>
* Partitions which do not exist in the metadata image are reported as unassigned.
*/
@Override
public boolean isPartitionAssigned(Uuid topicId, int partitionId) {
Map<Integer, String> partitionMap = invertedMemberAssignment.get(topicId);
if (partitionMap == null) {
return false;
if (assignedTopicPartitions == null) {
assignedTopicPartitions = computeAssignedTopicPartitions();
}
return partitionMap.containsKey(partitionId);
AssignedPartitions assignedPartitions = assignedTopicPartitions.get(topicId);
return assignedPartitions != null && assignedPartitions.partitions().get(partitionId);
}
private Map<Uuid, AssignedPartitions> computeAssignedTopicPartitions() {
Map<Uuid, AssignedPartitions> assignedTopicPartitions = new HashMap<>();
for (MemberSubscriptionAndAssignmentImpl member : members.values()) {
for (Map.Entry<Uuid, Set<Integer>> entry : member.partitions().entrySet()) {
AssignedPartitions assignedPartitions = assignedTopicPartitions.computeIfAbsent(
entry.getKey(),
topicId -> {
int partitionCount = metadataImage.topicMetadata(topicId)
.map(CoordinatorMetadataImage.TopicMetadata::partitionCount)
.orElse(0);
return new AssignedPartitions(new BitSet(partitionCount), partitionCount);
}
);
for (int partition : entry.getValue()) {
if (partition >= assignedPartitions.partitionCount()) {
// Drop partitions that do not exist.
continue;
}
assignedPartitions.partitions().set(partition);
}
}
}
return assignedTopicPartitions;
}
/**
@@ -141,13 +186,12 @@ public class GroupSpecImpl implements GroupSpec {
if (!(o instanceof GroupSpecImpl groupSpec)) return false;
return Objects.equals(members, groupSpec.members) &&
subscriptionType == groupSpec.subscriptionType &&
Objects.equals(invertedMemberAssignment, groupSpec.invertedMemberAssignment) &&
Objects.equals(topicPartitionAllowedMap, groupSpec.topicPartitionAllowedMap);
}
@Override
public int hashCode() {
return Objects.hash(members, subscriptionType, invertedMemberAssignment, topicPartitionAllowedMap);
return Objects.hash(members, subscriptionType, topicPartitionAllowedMap);
}
@Override
@@ -155,7 +199,6 @@ public class GroupSpecImpl implements GroupSpec {
return "GroupSpecImpl(" +
"members=" + members +
", subscriptionType=" + subscriptionType +
", invertedMemberAssignment=" + invertedMemberAssignment +
", topicPartitionAllowedMap=" + topicPartitionAllowedMap +
')';
}
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.coordinator.group.modern;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.errors.UnknownMemberIdException;
import org.apache.kafka.common.message.ListGroupsResponseData;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
@@ -32,7 +31,6 @@ import org.apache.kafka.timeline.TimelineObject;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
@@ -106,12 +104,6 @@ public abstract class ModernGroup<T extends ModernGroupMember> implements Group
*/
protected final TimelineHashMap<String, Assignment> targetAssignment;
/**
* Reverse lookup map representing topic partitions with
* their current member assignments.
*/
private final TimelineHashMap<Uuid, TimelineHashMap<Integer, String>> invertedTargetAssignment;
/**
* The metadata refresh deadline. It consists of a timestamp in milliseconds together with
* the group epoch at the time of setting it. The metadata refresh time is considered as a
@@ -138,7 +130,6 @@ public abstract class ModernGroup<T extends ModernGroupMember> implements Group
this.subscriptionType = new TimelineObject<>(snapshotRegistry, HOMOGENEOUS);
this.targetAssignmentMetadata = new TimelineObject<>(snapshotRegistry, TargetAssignmentMetadata.INITIAL);
this.targetAssignment = new TimelineHashMap<>(snapshotRegistry, 0);
this.invertedTargetAssignment = new TimelineHashMap<>(snapshotRegistry, 0);
}
/**
@@ -262,14 +253,6 @@ public abstract class ModernGroup<T extends ModernGroupMember> implements Group
return targetAssignment.getOrDefault(memberId, Assignment.EMPTY);
}
/**
* @return An immutable map containing all the topic partitions
* with their current member assignments.
*/
public Map<Uuid, Map<Integer, String>> invertedTargetAssignment() {
return Collections.unmodifiableMap(invertedTargetAssignment);
}
/**
* Updates target assignment of a member.
*
@@ -277,75 +260,15 @@ public abstract class ModernGroup<T extends ModernGroupMember> implements Group
* @param newTargetAssignment The new target assignment.
*/
public void updateTargetAssignment(String memberId, Assignment newTargetAssignment) {
updateInvertedTargetAssignment(
memberId,
targetAssignment.getOrDefault(memberId, new Assignment(Map.of())),
newTargetAssignment
);
targetAssignment.put(memberId, newTargetAssignment);
}
/**
* Updates the reverse lookup map of the target assignment.
*
* @param memberId The member Id.
* @param oldTargetAssignment The old target assignment.
* @param newTargetAssignment The new target assignment.
*/
private void updateInvertedTargetAssignment(
String memberId,
Assignment oldTargetAssignment,
Assignment newTargetAssignment
) {
// Combine keys from both old and new assignments.
Set<Uuid> allTopicIds = new HashSet<>();
allTopicIds.addAll(oldTargetAssignment.partitions().keySet());
allTopicIds.addAll(newTargetAssignment.partitions().keySet());
for (Uuid topicId : allTopicIds) {
Set<Integer> oldPartitions = oldTargetAssignment.partitions().getOrDefault(topicId, Set.of());
Set<Integer> newPartitions = newTargetAssignment.partitions().getOrDefault(topicId, Set.of());
TimelineHashMap<Integer, String> topicPartitionAssignment = invertedTargetAssignment.computeIfAbsent(
topicId, k -> new TimelineHashMap<>(snapshotRegistry, Math.max(oldPartitions.size(), newPartitions.size()))
);
// Remove partitions that aren't present in the new assignment only if the partition is currently
// still assigned to the member in question.
// If p0 was moved from A to B, and the target assignment map was updated for B first, we don't want to
// remove the key p0 from the inverted map and undo the action when A eventually tries to update its assignment.
for (Integer partition : oldPartitions) {
if (!newPartitions.contains(partition) && memberId.equals(topicPartitionAssignment.get(partition))) {
topicPartitionAssignment.remove(partition);
}
}
// Add partitions that are in the new assignment but not in the old assignment.
for (Integer partition : newPartitions) {
if (!oldPartitions.contains(partition)) {
topicPartitionAssignment.put(partition, memberId);
}
}
if (topicPartitionAssignment.isEmpty()) {
invertedTargetAssignment.remove(topicId);
} else {
invertedTargetAssignment.put(topicId, topicPartitionAssignment);
}
}
}
/**
* Removes the target assignment of a member.
*
* @param memberId The member id.
*/
public void removeTargetAssignment(String memberId) {
updateInvertedTargetAssignment(
memberId,
targetAssignment.getOrDefault(memberId, Assignment.EMPTY),
Assignment.EMPTY
);
targetAssignment.remove(memberId);
}
@@ -18,7 +18,6 @@ package org.apache.kafka.coordinator.group;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
import java.util.AbstractMap;
import java.util.Arrays;
@@ -108,32 +107,4 @@ public class AssignmentTestUtil {
assertEquals(expectedAssignment.get(memberId), computedAssignmentForMember);
});
}
/**
* Generate a reverse look up map of partition to member target assignments from the given metadata.
*
* @param members The member subscription specs.
* @return Map of topic partition to member assignments.
*/
public static Map<Uuid, Map<Integer, String>> invertedTargetAssignment(
Map<String, MemberSubscriptionAndAssignmentImpl> members
) {
Map<Uuid, Map<Integer, String>> invertedTargetAssignment = new HashMap<>();
for (Map.Entry<String, MemberSubscriptionAndAssignmentImpl> memberEntry : members.entrySet()) {
String memberId = memberEntry.getKey();
Map<Uuid, Set<Integer>> memberAssignment = memberEntry.getValue().partitions();
for (Map.Entry<Uuid, Set<Integer>> topicEntry : memberAssignment.entrySet()) {
Uuid topicId = topicEntry.getKey();
Set<Integer> partitions = topicEntry.getValue();
Map<Integer, String> partitionMap = invertedTargetAssignment.computeIfAbsent(topicId, k -> new HashMap<>());
for (Integer partitionId : partitions) {
partitionMap.put(partitionId, memberId);
}
}
}
return invertedTargetAssignment;
}
}
@@ -17,7 +17,7 @@
package org.apache.kafka.coordinator.group.assignor;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.common.runtime.KRaftCoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.MetadataImageBuilder;
import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
@@ -27,7 +27,6 @@ import org.apache.kafka.coordinator.group.modern.Assignment;
import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
import org.apache.kafka.image.MetadataImage;
import java.util.HashMap;
import java.util.LinkedHashMap;
@@ -37,7 +36,6 @@ import java.util.Optional;
import java.util.Set;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.assertAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.invertedTargetAssignment;
import static org.junit.jupiter.api.Assertions.assertSame;
public class CommonAssignorTests {
@@ -58,15 +56,15 @@ public class CommonAssignorTests {
* @param rackAware Whether to test with rack awareness.
*/
public static void testAssignmentReuse(PartitionAssignor assignor, SubscriptionType subscriptionType, boolean rackAware) {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 2)
.addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 5)
.addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 7)
.addRacks()
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = new HashMap<>();
@@ -90,9 +88,9 @@ public class CommonAssignorTests {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
subscriptionType,
Map.of()
subscriptionType
);
GroupAssignment firstAssignment = assignor.assign(
@@ -112,9 +110,9 @@ public class CommonAssignorTests {
));
}
GroupSpec groupSpecWithAssignment = new GroupSpecImpl(
metadataImage,
membersWithAssignment,
subscriptionType,
invertedTargetAssignment(membersWithAssignment)
subscriptionType
);
GroupAssignment secondAssignment = assignor.assign(
@@ -137,15 +135,15 @@ public class CommonAssignorTests {
* @param rackAware Whether to test with rack awareness.
*/
public static void testReassignmentStickiness(PartitionAssignor assignor, SubscriptionType subscriptionType, boolean rackAware) {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 2)
.addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 5)
.addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 7)
.addRacks()
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = new HashMap<>();
@@ -172,9 +170,9 @@ public class CommonAssignorTests {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
subscriptionType,
Map.of()
subscriptionType
);
GroupAssignment firstAssignment = assignor.assign(
@@ -209,9 +207,9 @@ public class CommonAssignorTests {
));
}
GroupSpec groupSpecWithAssignment = new GroupSpecImpl(
metadataImage,
membersWithAssignment,
subscriptionType,
invertedTargetAssignment(membersWithAssignment)
subscriptionType
);
GroupAssignment secondAssignment = assignor.assign(
@@ -17,6 +17,8 @@
package org.apache.kafka.coordinator.group.assignor;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.MetadataImageBuilder;
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;
@@ -38,18 +40,20 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
public class GroupSpecImplTest {
private static final String TEST_MEMBER = "test-member";
private CoordinatorMetadataImage metadataImage;
private Map<String, MemberSubscriptionAndAssignmentImpl> members;
private SubscriptionType subscriptionType;
private Map<Uuid, Map<Integer, String>> invertedTargetAssignment;
private GroupSpecImpl groupSpec;
private Uuid topicId;
@BeforeEach
void setUp() {
topicId = Uuid.randomUuid();
metadataImage = new MetadataImageBuilder()
.addTopic(topicId, "topic", 3)
.buildCoordinatorMetadataImage();
members = new HashMap<>();
subscriptionType = SubscriptionType.HOMOGENEOUS;
invertedTargetAssignment = new HashMap<>();
topicId = Uuid.randomUuid();
members.put(TEST_MEMBER, new MemberSubscriptionAndAssignmentImpl(
Optional.empty(),
@@ -59,9 +63,9 @@ public class GroupSpecImplTest {
));
groupSpec = new GroupSpecImpl(
metadataImage,
members,
subscriptionType,
invertedTargetAssignment
subscriptionType
);
}
@@ -77,9 +81,12 @@ public class GroupSpecImplTest {
@Test
void testIsPartitionAssigned() {
Map<Integer, String> partitionMap = new HashMap<>();
partitionMap.put(1, "test-member");
invertedTargetAssignment.put(topicId, partitionMap);
members.put(TEST_MEMBER, new MemberSubscriptionAndAssignmentImpl(
Optional.empty(),
Optional.empty(),
Set.of(topicId),
new Assignment(Map.of(topicId, Set.of(1)))
));
assertTrue(groupSpec.isPartitionAssigned(topicId, 1));
assertFalse(groupSpec.isPartitionAssigned(topicId, 2));
@@ -117,9 +124,9 @@ public class GroupSpecImplTest {
// Allowed map with data.
groupSpec = new GroupSpecImpl(
metadataImage,
members,
subscriptionType,
invertedTargetAssignment,
Optional.of(Map.of(topicId, Set.of(0)))
);
@@ -18,7 +18,6 @@ package org.apache.kafka.coordinator.group.assignor;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.KRaftCoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.MetadataImageBuilder;
import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
@@ -31,7 +30,6 @@ import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
import org.apache.kafka.coordinator.group.modern.MemberAssignmentImpl;
import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
import org.apache.kafka.image.MetadataImage;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
@@ -43,7 +41,6 @@ import java.util.Optional;
import java.util.Set;
import java.util.TreeMap;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.invertedTargetAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
import static org.apache.kafka.coordinator.group.api.assignor.SubscriptionType.HETEROGENEOUS;
@@ -76,8 +73,9 @@ public class RangeAssignorTest {
@Test
public void testOneMemberNoTopic() {
CoordinatorMetadataImage metadataImage = CoordinatorMetadataImage.EMPTY;
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
CoordinatorMetadataImage.EMPTY
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -91,9 +89,9 @@ public class RangeAssignorTest {
);
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
GroupAssignment groupAssignment = assignor.assign(
@@ -111,11 +109,11 @@ public class RangeAssignorTest {
@Test
public void testOneMemberSubscribedToNonExistentTopic() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -129,9 +127,9 @@ public class RangeAssignorTest {
);
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
assertThrows(PartitionAssignorException.class,
@@ -140,10 +138,10 @@ public class RangeAssignorTest {
@Test
public void testFirstAssignmentTwoMembersTwoTopicsSameSubscriptions() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic3Uuid, topic3Name, 2)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -162,12 +160,12 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -190,11 +188,11 @@ public class RangeAssignorTest {
@Test
public void testFirstAssignmentThreeMembersThreeTopicsDifferentSubscriptions() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic2Uuid, topic2Name, 3)
.addTopic(topic3Uuid, topic3Name, 2)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -220,12 +218,12 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -251,10 +249,10 @@ public class RangeAssignorTest {
@Test
public void testFirstAssignmentNumMembersGreaterThanNumPartitions() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic3Uuid, topic3Name, 2)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -280,12 +278,12 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -312,11 +310,11 @@ public class RangeAssignorTest {
@Test
public void testStaticMembership() throws PartitionAssignorException {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriber subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -334,9 +332,9 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
SubscriptionType.HOMOGENEOUS,
invertedTargetAssignment(members)
SubscriptionType.HOMOGENEOUS
);
GroupAssignment initialAssignment = assignor.assign(
@@ -354,9 +352,9 @@ public class RangeAssignorTest {
));
groupSpec = new GroupSpecImpl(
metadataImage,
members,
SubscriptionType.HOMOGENEOUS,
invertedTargetAssignment(members)
SubscriptionType.HOMOGENEOUS
);
GroupAssignment reassignedAssignment = assignor.assign(
@@ -378,11 +376,11 @@ public class RangeAssignorTest {
@Test
public void testMixedStaticMembership() throws PartitionAssignorException {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 5)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriber subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
// Initialize members with instance Ids.
@@ -409,9 +407,9 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
SubscriptionType.HOMOGENEOUS,
invertedTargetAssignment(members)
SubscriptionType.HOMOGENEOUS
);
GroupAssignment initialAssignment = assignor.assign(
@@ -429,9 +427,9 @@ public class RangeAssignorTest {
));
groupSpec = new GroupSpecImpl(
metadataImage,
members,
SubscriptionType.HOMOGENEOUS,
invertedTargetAssignment(members)
SubscriptionType.HOMOGENEOUS
);
GroupAssignment reassignedAssignment = assignor.assign(
@@ -458,10 +456,10 @@ public class RangeAssignorTest {
@Test
public void testReassignmentNumMembersGreaterThanNumPartitionsWhenOneMemberAdded() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 2)
.addTopic(topic2Uuid, topic2Name, 2)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -494,12 +492,12 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -525,10 +523,10 @@ public class RangeAssignorTest {
@Test
public void testReassignmentWhenOnePartitionAddedForTwoMembersTwoTopics() {
// Simulating adding a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 4)
.addTopic(topic2Uuid, topic2Name, 4)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -553,12 +551,12 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -581,10 +579,10 @@ public class RangeAssignorTest {
@Test
public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembersTwoTopics() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic2Uuid, topic2Name, 3)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -617,12 +615,12 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -650,10 +648,10 @@ public class RangeAssignorTest {
@Test
public void testReassignmentWhenOneMemberAddedAndOnePartitionAfterInitialAssignmentWithTwoMembersTwoTopics() {
// Add a new partition to topic 1, initially T1 -> 3 partitions
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 4)
.addTopic(topic2Uuid, topic2Name, 3)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -686,12 +684,12 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -717,10 +715,10 @@ public class RangeAssignorTest {
@Test
public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithTwoMembersTwoTopics() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic2Uuid, topic2Name, 3)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -737,12 +735,12 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -761,11 +759,11 @@ public class RangeAssignorTest {
@Test
public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignmentWithThreeMembersTwoTopics() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic2Uuid, topic2Name, 3)
.addTopic(topic3Uuid, topic3Name, 2)
.build();
.buildCoordinatorMetadataImage();
// Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3
// Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2
@@ -801,12 +799,12 @@ public class RangeAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -18,7 +18,6 @@ package org.apache.kafka.coordinator.group.assignor;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.KRaftCoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.MetadataImageBuilder;
import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
@@ -28,7 +27,6 @@ import org.apache.kafka.coordinator.group.modern.Assignment;
import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
import org.apache.kafka.image.MetadataImage;
import org.apache.kafka.server.common.TopicIdPartition;
import org.junit.jupiter.api.Test;
@@ -71,14 +69,15 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithEmptyMembers() {
CoordinatorMetadataImage metadataImage = CoordinatorMetadataImage.EMPTY;
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
CoordinatorMetadataImage.EMPTY
metadataImage
);
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
Map.of(),
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
GroupAssignment groupAssignment = assignor.assign(
@@ -89,9 +88,9 @@ public class SimpleAssignorTest {
assertEquals(Map.of(), groupAssignment.members());
groupSpec = new GroupSpecImpl(
metadataImage,
Map.of(),
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
);
groupAssignment = assignor.assign(
groupSpec,
@@ -102,11 +101,11 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithNoSubscribedTopic() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -120,9 +119,9 @@ public class SimpleAssignorTest {
);
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
GroupAssignment groupAssignment = assignor.assign(
@@ -135,11 +134,11 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithSubscribedToNonExistentTopic() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -153,9 +152,9 @@ public class SimpleAssignorTest {
);
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
assertThrows(PartitionAssignorException.class,
@@ -164,10 +163,10 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithTwoMembersAndTwoTopicsHomogeneous() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new HashMap<>();
@@ -190,12 +189,12 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -208,10 +207,10 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithTwoMembersAndTwoTopicsHomogeneousWithAllowedMap() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 3)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new HashMap<>();
@@ -234,9 +233,9 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of(),
Optional.of(
Map.of(
TOPIC_1_UUID, Set.of(0, 1, 2),
@@ -245,7 +244,7 @@ public class SimpleAssignorTest {
)
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -258,10 +257,10 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithTwoMembersAndTwoTopicsHomogeneousWithNonAssignableTopic() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new HashMap<>();
@@ -284,16 +283,16 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of(),
Optional.of(
Map.of(TOPIC_1_UUID, Set.of(0, 1, 2))
)
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -306,11 +305,11 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithThreeMembersThreeTopicsHeterogeneous() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3)
.addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2)
.build();
.buildCoordinatorMetadataImage();
Set<Uuid> memberATopicsSubscription = new LinkedHashSet<>();
memberATopicsSubscription.add(TOPIC_1_UUID);
@@ -342,12 +341,12 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -361,11 +360,11 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithThreeMembersThreeTopicsHeterogeneousWithAllowedMap() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3)
.addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2)
.build();
.buildCoordinatorMetadataImage();
Set<Uuid> memberATopicsSubscription = new LinkedHashSet<>();
memberATopicsSubscription.add(TOPIC_1_UUID);
@@ -397,9 +396,9 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
Map.of(),
Optional.of(
Map.of(
TOPIC_1_UUID, Set.of(0, 1), // but not 2
@@ -410,7 +409,7 @@ public class SimpleAssignorTest {
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -424,11 +423,11 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithThreeMembersThreeTopicsHeterogeneousWithNonAssignableTopic() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3)
.addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2) // non-assignable
.build();
.buildCoordinatorMetadataImage();
Set<Uuid> memberATopicsSubscription = new LinkedHashSet<>();
memberATopicsSubscription.add(TOPIC_1_UUID);
@@ -460,9 +459,9 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
Map.of(),
Optional.of(
Map.of(
TOPIC_1_UUID, Set.of(0, 1, 2),
@@ -472,7 +471,7 @@ public class SimpleAssignorTest {
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -497,10 +496,10 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithOneMemberNoAssignedTopicHeterogeneous() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 2)
.build();
.buildCoordinatorMetadataImage();
Set<Uuid> memberATopicsSubscription = new LinkedHashSet<>();
memberATopicsSubscription.add(TOPIC_1_UUID);
@@ -521,12 +520,12 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -550,12 +549,12 @@ public class SimpleAssignorTest {
final int numPartitions = 24;
final int numMembers = 101;
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, numPartitions)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Set<Uuid> topicsSubscription = new LinkedHashSet<>();
@@ -575,9 +574,9 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
new HashMap<>()
HOMOGENEOUS
);
GroupAssignment computedAssignment = assignor.assign(groupSpec, subscribedTopicMetadata);
@@ -597,12 +596,12 @@ public class SimpleAssignorTest {
final int numPartitions = 24;
final int numMembers = 101;
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, numPartitions)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Set<Uuid> topicsSubscription = new LinkedHashSet<>();
@@ -622,9 +621,9 @@ public class SimpleAssignorTest {
}
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
new HashMap<>()
HOMOGENEOUS
);
GroupAssignment computedAssignment = assignor.assign(groupSpec, subscribedTopicMetadata);
@@ -646,9 +645,9 @@ public class SimpleAssignorTest {
members.remove(newMemberId);
groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
new HashMap<>()
HOMOGENEOUS
);
computedAssignment = assignor.assign(groupSpec, subscribedTopicMetadata);
@@ -666,11 +665,11 @@ public class SimpleAssignorTest {
@Test
public void testAssignWithCurrentAssignmentHeterogeneous() {
// Current assignment setup - 3 members A - {T1, T2}, B - {T3}, C - {T2, T3}.
MetadataImage metadataImage1 = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage1 = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3)
.addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2)
.build();
.buildCoordinatorMetadataImage();
Set<Uuid> memberATopicsSubscription1 = new LinkedHashSet<>();
memberATopicsSubscription1.add(TOPIC_1_UUID);
@@ -702,12 +701,12 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec1 = new GroupSpecImpl(
metadataImage1,
members1,
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage1)
metadataImage1
);
GroupAssignment computedAssignment1 = assignor.assign(
@@ -719,12 +718,12 @@ public class SimpleAssignorTest {
// New assignment setup - 2 members A - {T1, T2, T3}, B - {T3, T4}.
MetadataImage metadataImage2 = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage2 = new MetadataImageBuilder()
.addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
.addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3)
.addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2)
.addTopic(TOPIC_4_UUID, TOPIC_4_NAME, 1)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members2 = new HashMap<>();
@@ -755,13 +754,13 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec2 = new GroupSpecImpl(
metadataImage2,
members2,
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage2)
metadataImage2
);
GroupAssignment computedAssignment2 = assignor.assign(
@@ -820,9 +819,9 @@ public class SimpleAssignorTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
new HashMap<>()
HETEROGENEOUS
);
GroupAssignment computedAssignment = assignor.assign(groupSpec, subscribedTopicMetadata);
@@ -888,9 +887,9 @@ public class SimpleAssignorTest {
}
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
new HashMap<>()
HETEROGENEOUS
);
GroupAssignment computedAssignment = assignor.assign(groupSpec, subscribedTopicMetadata);
@@ -912,9 +911,9 @@ public class SimpleAssignorTest {
members.remove(newMemberId);
groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
new HashMap<>()
HETEROGENEOUS
);
computedAssignment = assignor.assign(groupSpec, subscribedTopicMetadata);
@@ -17,7 +17,7 @@
package org.apache.kafka.coordinator.group.assignor;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.common.runtime.KRaftCoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.MetadataImageBuilder;
import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
@@ -26,7 +26,6 @@ import org.apache.kafka.coordinator.group.modern.Assignment;
import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
import org.apache.kafka.image.MetadataImage;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
@@ -39,7 +38,6 @@ import java.util.Set;
import java.util.TreeMap;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.assertAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.invertedTargetAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
import static org.apache.kafka.coordinator.group.api.assignor.SubscriptionType.HETEROGENEOUS;
@@ -74,11 +72,11 @@ public class UniformHeterogeneousAssignmentBuilderTest {
@Test
public void testTwoMembersNoTopicSubscription() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -96,9 +94,9 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
);
GroupAssignment groupAssignment = assignor.assign(
@@ -111,11 +109,11 @@ public class UniformHeterogeneousAssignmentBuilderTest {
@Test
public void testTwoMembersSubscribedToNonexistentTopics() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -133,9 +131,9 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
);
assertThrows(
@@ -145,10 +143,10 @@ public class UniformHeterogeneousAssignmentBuilderTest {
@Test
public void testFirstAssignmentTwoMembersTwoTopics() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic3Uuid, topic3Name, 6)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -167,12 +165,12 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -194,10 +192,10 @@ public class UniformHeterogeneousAssignmentBuilderTest {
@Test
public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 2)
.addTopic(topic3Uuid, topic3Name, 1)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -223,12 +221,12 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -253,11 +251,11 @@ public class UniformHeterogeneousAssignmentBuilderTest {
@Test
public void testReassignmentForTwoMembersThreeTopicsGivenUnbalancedPrevAssignment() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 6)
.addTopic(topic2Uuid, topic2Name, 4)
.addTopic(topic3Uuid, topic3Name, 4)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -292,12 +290,12 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -324,12 +322,12 @@ public class UniformHeterogeneousAssignmentBuilderTest {
@Test
public void testReassignmentWhenPartitionsAreAddedForTwoMembers() {
// Simulating adding partitions to T1, T2, T3 - originally T1 -> 4, T2 -> 3, T3 -> 2, T4 -> 3
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 6)
.addTopic(topic2Uuid, topic2Name, 5)
.addTopic(topic3Uuid, topic3Name, 3)
.addTopic(topic4Uuid, topic4Name, 3)
.build();
.buildCoordinatorMetadataImage();
@@ -356,12 +354,12 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -385,10 +383,10 @@ public class UniformHeterogeneousAssignmentBuilderTest {
@Test
public void testReassignmentWhenOneMemberAddedAndPartitionsAddedTwoMembersTwoTopics() {
// Initially T1 -> 3, T2 -> 3 partitions.
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 6)
.addTopic(topic2Uuid, topic2Name, 7)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -421,12 +419,12 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -451,11 +449,11 @@ public class UniformHeterogeneousAssignmentBuilderTest {
@Test
public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeMembersThreeTopics() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic2Uuid, topic2Name, 8)
.addTopic(topic3Uuid, topic3Name, 3)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -481,12 +479,12 @@ public class UniformHeterogeneousAssignmentBuilderTest {
// Member C was removed
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -508,10 +506,10 @@ public class UniformHeterogeneousAssignmentBuilderTest {
@Test
public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWithTwoMembersTwoTopics() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic2Uuid, topic2Name, 5)
.build();
.buildCoordinatorMetadataImage();
// Initial subscriptions were [T1, T2]
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -537,12 +535,12 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -569,11 +567,11 @@ public class UniformHeterogeneousAssignmentBuilderTest {
*/
@Test
public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 2)
.addTopic(topic2Uuid, topic2Name, 2)
.addTopic(topic3Uuid, topic3Name, 2)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -596,12 +594,12 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
invertedTargetAssignment(members)
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -624,9 +622,9 @@ public class UniformHeterogeneousAssignmentBuilderTest {
@Test
public void testFirstAssignmentWithTwoMembersIncludingOneWithoutSubscriptions() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -645,12 +643,12 @@ public class UniformHeterogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -17,7 +17,7 @@
package org.apache.kafka.coordinator.group.assignor;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.common.runtime.KRaftCoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.MetadataImageBuilder;
import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
@@ -26,7 +26,6 @@ import org.apache.kafka.coordinator.group.modern.Assignment;
import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
import org.apache.kafka.image.MetadataImage;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
@@ -42,7 +41,6 @@ import java.util.Set;
import java.util.TreeMap;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.assertAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.invertedTargetAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkOrderedAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
@@ -77,11 +75,11 @@ public class UniformHomogeneousAssignmentBuilderTest {
@Test
public void testOneMemberNoTopicSubscription() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -95,9 +93,9 @@ public class UniformHomogeneousAssignmentBuilderTest {
);
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
GroupAssignment groupAssignment = assignor.assign(
@@ -110,11 +108,11 @@ public class UniformHomogeneousAssignmentBuilderTest {
@Test
public void testOneMemberSubscribedToNonexistentTopic() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.build();
.buildCoordinatorMetadataImage();
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -128,9 +126,9 @@ public class UniformHomogeneousAssignmentBuilderTest {
);
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
assertThrows(PartitionAssignorException.class,
@@ -139,10 +137,10 @@ public class UniformHomogeneousAssignmentBuilderTest {
@Test
public void testFirstAssignmentTwoMembersTwoTopicsNoMemberRacks() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic3Uuid, topic3Name, 2)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -169,12 +167,12 @@ public class UniformHomogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -188,9 +186,9 @@ public class UniformHomogeneousAssignmentBuilderTest {
@Test
public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic3Uuid, topic3Name, 2)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -228,12 +226,12 @@ public class UniformHomogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -254,7 +252,7 @@ public class UniformHomogeneousAssignmentBuilderTest {
metadataImageBuilder.addTopic(topicId, "topic-" + i, 3);
topicIds.add(topicId);
}
MetadataImage metadataImage = metadataImageBuilder.build();
CoordinatorMetadataImage metadataImage = metadataImageBuilder.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
for (int i = 1; i < 50; i++) {
@@ -267,12 +265,12 @@ public class UniformHomogeneousAssignmentBuilderTest {
}
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -285,10 +283,10 @@ public class UniformHomogeneousAssignmentBuilderTest {
@Test
public void testReassignmentForTwoMembersTwoTopicsGivenUnbalancedPrevAssignment() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic2Uuid, topic2Name, 3)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -323,12 +321,12 @@ public class UniformHomogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -343,10 +341,10 @@ public class UniformHomogeneousAssignmentBuilderTest {
@Test
public void testReassignmentWhenPartitionsAreAddedForTwoMembersTwoTopics() {
// Simulating adding partition to T1 and T2 - originally T1 -> 3 Partitions and T2 -> 3 Partitions
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 6)
.addTopic(topic2Uuid, topic2Name, 5)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -381,12 +379,12 @@ public class UniformHomogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -400,10 +398,10 @@ public class UniformHomogeneousAssignmentBuilderTest {
@Test
public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembersTwoTopics() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic2Uuid, topic2Name, 3)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -448,12 +446,12 @@ public class UniformHomogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -467,10 +465,10 @@ public class UniformHomogeneousAssignmentBuilderTest {
@Test
public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeMembersTwoTopics() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 3)
.addTopic(topic2Uuid, topic2Name, 3)
.build();
.buildCoordinatorMetadataImage();
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -507,12 +505,12 @@ public class UniformHomogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -526,10 +524,10 @@ public class UniformHomogeneousAssignmentBuilderTest {
@Test
public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWithTwoMembersTwoTopics() {
MetadataImage metadataImage = new MetadataImageBuilder()
CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
.addTopic(topic1Uuid, topic1Name, 2)
.addTopic(topic2Uuid, topic2Name, 2)
.build();
.buildCoordinatorMetadataImage();
// Initial subscriptions were [T1, T2]
Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -563,12 +561,12 @@ public class UniformHomogeneousAssignmentBuilderTest {
));
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
members,
HOMOGENEOUS,
invertedTargetAssignment(members)
HOMOGENEOUS
);
SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
new KRaftCoordinatorMetadataImage(metadataImage)
metadataImage
);
GroupAssignment computedAssignment = assignor.assign(
@@ -17,6 +17,7 @@
package org.apache.kafka.coordinator.group.assignor.uniform2;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType;
import org.apache.kafka.coordinator.group.modern.Assignment;
@@ -99,7 +100,8 @@ public final class GroupSpecFixture {
checkNotBuilt();
built = true;
var subscriptionType = sameSubscriptions ? SubscriptionType.HOMOGENEOUS : SubscriptionType.HETEROGENEOUS;
return new GroupSpecImpl(members, subscriptionType, invertedTargetAssignment);
// The metadata image is only used by isPartitionAssigned, which the uniform2 assignor never calls.
return new GroupSpecImpl(CoordinatorMetadataImage.EMPTY, members, subscriptionType);
}
private void checkNotBuilt() {
@@ -51,15 +51,14 @@ public class GroupSpecBuilderTest {
.withMembers(Map.of())
.withSubscriptionType(HOMOGENEOUS)
.withTargetAssignment(Map.of())
.withInvertedTargetAssignment(Map.of())
.withMetadataImage(metadataImage)
.withResolvedRegularExpressions(Map.of());
assertEquals(
new GroupSpecImpl(
metadataImage,
Map.of(),
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
),
builder.build()
);
@@ -111,29 +110,12 @@ public class GroupSpecBuilderTest {
"member-2", assignment2,
"member-3", assignment3
))
.withInvertedTargetAssignment(Map.of(
fooTopicId, Map.of(
1, "member-1",
2, "member-1",
3, "member-2",
4, "member-2",
5, "member-3",
6, "member-3"
),
barTopicId, Map.of(
1, "member-1",
2, "member-1",
3, "member-2",
4, "member-2",
5, "member-3",
6, "member-3"
)
))
.withMetadataImage(metadataImage)
.withResolvedRegularExpressions(Map.of());
assertEquals(
new GroupSpecImpl(
metadataImage,
Map.of(
"member-1", new MemberSubscriptionAndAssignmentImpl(
Optional.empty(),
@@ -154,25 +136,7 @@ public class GroupSpecBuilderTest {
assignment3
)
),
HETEROGENEOUS,
Map.of(
fooTopicId, Map.of(
1, "member-1",
2, "member-1",
3, "member-2",
4, "member-2",
5, "member-3",
6, "member-3"
),
barTopicId, Map.of(
1, "member-1",
2, "member-1",
3, "member-2",
4, "member-2",
5, "member-3",
6, "member-3"
)
)
HETEROGENEOUS
),
builder.build()
);
@@ -209,7 +173,6 @@ public class GroupSpecBuilderTest {
"member-2", Assignment.EMPTY,
"member-3", Assignment.EMPTY
))
.withInvertedTargetAssignment(Map.of())
.withMetadataImage(metadataImage)
.withResolvedRegularExpressions(Map.of(
"foo*", new ResolvedRegularExpression(
@@ -221,6 +184,7 @@ public class GroupSpecBuilderTest {
assertEquals(
new GroupSpecImpl(
metadataImage,
Map.of(
"member-1", new MemberSubscriptionAndAssignmentImpl(
Optional.empty(),
@@ -241,8 +205,7 @@ public class GroupSpecBuilderTest {
Assignment.EMPTY
)
),
HETEROGENEOUS,
Map.of()
HETEROGENEOUS
),
builder.build()
);
@@ -268,14 +231,6 @@ public class GroupSpecBuilderTest {
))
));
Map<Integer, String> fooInvertedTargetAssignment = new HashMap<>(Map.of(
1, "member-1",
2, "member-1"
));
Map<Uuid, Map<Integer, String>> invertedTargetAssignment = new HashMap<>(Map.of(
fooTopicId, fooInvertedTargetAssignment
));
Map<String, ResolvedRegularExpression> resolvedRegularExpressions = new HashMap<>(Map.of(
"foo*", new ResolvedRegularExpression(
Set.of("foo"),
@@ -293,7 +248,6 @@ public class GroupSpecBuilderTest {
.withMembers(members)
.withSubscriptionType(HOMOGENEOUS)
.withTargetAssignment(targetAssignment)
.withInvertedTargetAssignment(invertedTargetAssignment)
.withMetadataImage(metadataImage)
.withResolvedRegularExpressions(resolvedRegularExpressions)
.withTopicAssignablePartitionsMap(topicAssignablePartitions)
@@ -303,14 +257,13 @@ public class GroupSpecBuilderTest {
// Modifications after the GroupSpec has been built should not be visible in the GroupSpec.
members.clear();
targetAssignment.clear();
fooInvertedTargetAssignment.clear();
invertedTargetAssignment.clear();
resolvedRegularExpressions.clear();
fooAssignablePartitions.clear();
topicAssignablePartitions.clear();
assertEquals(
new GroupSpecImpl(
metadataImage,
// members
Map.of("member-1", new MemberSubscriptionAndAssignmentImpl(
Optional.empty(),
@@ -324,8 +277,6 @@ public class GroupSpecBuilderTest {
new Assignment(mkAssignment(mkTopicAssignment(fooTopicId, 1, 2)))
)),
HOMOGENEOUS,
// invertedTargetAssignment
Map.of(fooTopicId, Map.of(1, "member-1", 2, "member-1")),
// topicAssignablePartitionsMap
Optional.of(Map.of(fooTopicId, Set.of(1, 2, 3)))
),
@@ -51,9 +51,9 @@ public class TargetAssignmentBuilderTest {
.buildCoordinatorMetadataImage();
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
Map.of(),
HOMOGENEOUS,
Map.of()
HOMOGENEOUS
);
GroupAssignment groupAssignment = new GroupAssignment(Map.of());
@@ -93,6 +93,7 @@ public class TargetAssignmentBuilderTest {
.buildCoordinatorMetadataImage();
GroupSpec groupSpec = new GroupSpecImpl(
metadataImage,
Map.of(
"member-1", new MemberSubscriptionAndAssignmentImpl(
Optional.empty(),
@@ -113,25 +114,7 @@ public class TargetAssignmentBuilderTest {
Assignment.EMPTY
)
),
HETEROGENEOUS,
Map.of(
fooTopicId, Map.of(
1, "member-1",
2, "member-1",
3, "member-2",
4, "member-2",
5, "member-3",
6, "member-3"
),
barTopicId, Map.of(
1, "member-1",
2, "member-1",
3, "member-2",
4, "member-2",
5, "member-3",
6, "member-3"
)
)
HETEROGENEOUS
);
GroupAssignment groupAssignment = new GroupAssignment(Map.of(
@@ -69,8 +69,6 @@ import java.util.OptionalLong;
import java.util.Set;
import java.util.stream.Stream;
import static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;
import static org.apache.kafka.coordinator.group.Assertions.assertUnorderedRecordsEquals;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignmentWithEpochs;
@@ -725,89 +723,6 @@ public class ConsumerGroupTest {
);
}
@Test
public void testUpdateInvertedAssignment() {
SnapshotRegistry snapshotRegistry = new SnapshotRegistry(new LogContext());
ConsumerGroup consumerGroup = new ConsumerGroup(new LogContext(), snapshotRegistry, "test-group");
Uuid topicId = Uuid.randomUuid();
String memberId1 = "member1";
String memberId2 = "member2";
// Initial assignment for member1
Assignment initialAssignment = new Assignment(Map.of(
topicId,
Set.of(0)
));
consumerGroup.updateTargetAssignment(memberId1, initialAssignment);
// Verify that partition 0 is assigned to member1.
assertEquals(
mkMap(
mkEntry(topicId, mkMap(mkEntry(0, memberId1)))
),
consumerGroup.invertedTargetAssignment()
);
// New assignment for member1
Assignment newAssignment = new Assignment(Map.of(
topicId,
Set.of(1)
));
consumerGroup.updateTargetAssignment(memberId1, newAssignment);
// Verify that partition 0 is no longer assigned and partition 1 is assigned to member1
assertEquals(
mkMap(
mkEntry(topicId, mkMap(mkEntry(1, memberId1)))
),
consumerGroup.invertedTargetAssignment()
);
// New assignment for member2 to add partition 1
Assignment newAssignment2 = new Assignment(Map.of(
topicId,
Set.of(1)
));
consumerGroup.updateTargetAssignment(memberId2, newAssignment2);
// Verify that partition 1 is assigned to member2
assertEquals(
mkMap(
mkEntry(topicId, mkMap(mkEntry(1, memberId2)))
),
consumerGroup.invertedTargetAssignment()
);
// New assignment for member1 to revoke partition 1 and assign partition 0
Assignment newAssignment1 = new Assignment(Map.of(
topicId,
Set.of(0)
));
consumerGroup.updateTargetAssignment(memberId1, newAssignment1);
// Verify that partition 1 is still assigned to member2 and partition 0 is assigned to member1
assertEquals(
mkMap(
mkEntry(topicId, mkMap(
mkEntry(0, memberId1),
mkEntry(1, memberId2)
))
),
consumerGroup.invertedTargetAssignment()
);
// Test remove target assignment for member1
consumerGroup.removeTargetAssignment(memberId1);
// Verify that partition 0 is no longer assigned and partition 1 is still assigned to member2
assertEquals(
mkMap(
mkEntry(topicId, mkMap(mkEntry(1, memberId2)))
),
consumerGroup.invertedTargetAssignment()
);
}
@Test
public void testMetadataRefreshDeadline() {
MockTime time = new MockTime();
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.coordinator.group.modern.share;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.errors.GroupIdNotFoundException;
import org.apache.kafka.common.errors.GroupNotEmptyException;
import org.apache.kafka.common.errors.UnknownMemberIdException;
@@ -28,7 +27,6 @@ import org.apache.kafka.common.utils.internals.LogContext;
import org.apache.kafka.coordinator.common.runtime.KRaftCoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.MetadataImageBuilder;
import org.apache.kafka.coordinator.group.Group;
import org.apache.kafka.coordinator.group.modern.Assignment;
import org.apache.kafka.coordinator.group.modern.MemberState;
import org.apache.kafka.coordinator.group.modern.share.ShareGroup.InitMapValue;
import org.apache.kafka.coordinator.group.modern.share.ShareGroup.ShareGroupState;
@@ -42,8 +40,6 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;
import static org.apache.kafka.coordinator.group.api.assignor.SubscriptionType.HETEROGENEOUS;
import static org.apache.kafka.coordinator.group.api.assignor.SubscriptionType.HOMOGENEOUS;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
@@ -222,89 +218,6 @@ public class ShareGroupTest {
);
}
@Test
public void testUpdateInvertedAssignment() {
SnapshotRegistry snapshotRegistry = new SnapshotRegistry(new LogContext());
ShareGroup shareGroup = new ShareGroup(snapshotRegistry, "test-group");
Uuid topicId = Uuid.randomUuid();
String memberId1 = "member1";
String memberId2 = "member2";
// Initial assignment for member1
Assignment initialAssignment = new Assignment(Map.of(
topicId,
Set.of(0)
));
shareGroup.updateTargetAssignment(memberId1, initialAssignment);
// Verify that partition 0 is assigned to member1.
assertEquals(
mkMap(
mkEntry(topicId, mkMap(mkEntry(0, memberId1)))
),
shareGroup.invertedTargetAssignment()
);
// New assignment for member1
Assignment newAssignment = new Assignment(Map.of(
topicId,
Set.of(1)
));
shareGroup.updateTargetAssignment(memberId1, newAssignment);
// Verify that partition 0 is no longer assigned and partition 1 is assigned to member1
assertEquals(
mkMap(
mkEntry(topicId, mkMap(mkEntry(1, memberId1)))
),
shareGroup.invertedTargetAssignment()
);
// New assignment for member2 to add partition 1
Assignment newAssignment2 = new Assignment(Map.of(
topicId,
Set.of(1)
));
shareGroup.updateTargetAssignment(memberId2, newAssignment2);
// Verify that partition 1 is assigned to member2
assertEquals(
mkMap(
mkEntry(topicId, mkMap(mkEntry(1, memberId2)))
),
shareGroup.invertedTargetAssignment()
);
// New assignment for member1 to revoke partition 1 and assign partition 0
Assignment newAssignment1 = new Assignment(Map.of(
topicId,
Set.of(0)
));
shareGroup.updateTargetAssignment(memberId1, newAssignment1);
// Verify that partition 1 is still assigned to member2 and partition 0 is assigned to member1
assertEquals(
mkMap(
mkEntry(topicId, mkMap(
mkEntry(0, memberId1),
mkEntry(1, memberId2)
))
),
shareGroup.invertedTargetAssignment()
);
// Test remove target assignment for member1
shareGroup.removeTargetAssignment(memberId1);
// Verify that partition 0 is no longer assigned and partition 1 is still assigned to member2
assertEquals(
mkMap(
mkEntry(topicId, mkMap(mkEntry(1, memberId2)))
),
shareGroup.invertedTargetAssignment()
);
}
@Test
public void testMetadataRefreshDeadline() {
MockTime time = new MockTime();
@@ -22,9 +22,7 @@ import org.apache.kafka.common.metadata.TopicRecord;
import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
import org.apache.kafka.coordinator.common.runtime.KRaftCoordinatorMetadataImage;
import org.apache.kafka.coordinator.group.Utils;
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.MemberAssignment;
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;
@@ -41,39 +39,10 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
public class AssignorBenchmarkUtils {
/**
* Generate a reverse look up map of partition to member target assignments from the given member spec.
*
* @param groupAssignment The group assignment.
* @return Map of topic partition to member assignments.
*/
public static Map<Uuid, Map<Integer, String>> computeInvertedTargetAssignment(
GroupAssignment groupAssignment
) {
Map<Uuid, Map<Integer, String>> invertedTargetAssignment = new HashMap<>();
for (Map.Entry<String, MemberAssignment> memberEntry : groupAssignment.members().entrySet()) {
String memberId = memberEntry.getKey();
Map<Uuid, Set<Integer>> topicsAndPartitions = memberEntry.getValue().partitions();
for (Map.Entry<Uuid, Set<Integer>> topicEntry : topicsAndPartitions.entrySet()) {
Uuid topicId = topicEntry.getKey();
Set<Integer> partitions = topicEntry.getValue();
Map<Integer, String> partitionMap = invertedTargetAssignment.computeIfAbsent(topicId, k -> new HashMap<>());
for (Integer partitionId : partitions) {
partitionMap.put(partitionId, memberId);
}
}
}
return invertedTargetAssignment;
}
/**
* Generates a list of topic names for use in benchmarks.
*
@@ -119,12 +88,14 @@ public class AssignorBenchmarkUtils {
/**
* Creates a GroupSpec from the given ConsumerGroupMembers.
*
* @param metadataImage The metadata image.
* @param members The ConsumerGroupMembers.
* @param subscriptionType The group's subscription type.
* @param topicResolver The TopicResolver to use.
* @return The new GroupSpec.
*/
public static GroupSpec createConsumerGroupSpec(
CoordinatorMetadataImage metadataImage,
Map<String, ConsumerGroupMember> members,
SubscriptionType subscriptionType,
TopicIds.TopicResolver topicResolver
@@ -144,21 +115,23 @@ public class AssignorBenchmarkUtils {
}
return new GroupSpecImpl(
metadataImage,
memberSpecs,
subscriptionType,
Map.of()
subscriptionType
);
}
/**
* Creates a GroupSpec from the given ShareGroupMembers.
*
* @param metadataImage The metadata image.
* @param members The ShareGroupMembers.
* @param subscriptionType The group's subscription type.
* @param topicResolver The TopicResolver to use.
* @return The new GroupSpec.
*/
public static GroupSpec createShareGroupSpec(
CoordinatorMetadataImage metadataImage,
Map<String, ShareGroupMember> members,
SubscriptionType subscriptionType,
TopicIds.TopicResolver topicResolver
@@ -178,9 +151,9 @@ public class AssignorBenchmarkUtils {
}
return new GroupSpecImpl(
metadataImage,
memberSpecs,
subscriptionType,
Map.of()
subscriptionType
);
}
@@ -390,7 +390,6 @@ public class ConsumerAssignorBenchmark {
}
var members = new HashMap<String, MemberSubscriptionAndAssignmentImpl>();
var invertedTargetAssignment = new HashMap<Uuid, Map<Integer, String>>();
for (int i = 0; i < memberCount; i++) {
var memberId = "member" + i;
var memberAssignment = currentAssignment.members().get(memberId);
@@ -402,8 +401,6 @@ public class ConsumerAssignorBenchmark {
var partitions = new HashMap<Uuid, Set<Integer>>();
memberAssignment.partitions().forEach((topicId, topicPartitions) -> {
partitions.put(topicId, new HashSet<>(topicPartitions));
var owners = invertedTargetAssignment.computeIfAbsent(topicId, id -> new HashMap<>());
topicPartitions.forEach(partition -> owners.put(partition, memberId));
});
assignment = new Assignment(partitions);
}
@@ -416,9 +413,9 @@ public class ConsumerAssignorBenchmark {
}
var spec = new GroupSpecImpl(
image,
members,
subscription == Subscription.HOMOGENEOUS ? SubscriptionType.HOMOGENEOUS : SubscriptionType.HETEROGENEOUS,
invertedTargetAssignment
subscription == Subscription.HOMOGENEOUS ? SubscriptionType.HOMOGENEOUS : SubscriptionType.HETEROGENEOUS
);
return new Group(spec, topicResolver, describer);
@@ -115,6 +115,8 @@ public class ShareGroupAssignorBenchmark {
private List<String> allTopicNames = List.of();
private CoordinatorMetadataImage metadataImage;
private TopicIds.TopicResolver topicResolver;
private SubscribedTopicDescriber subscribedTopicDescriber;
@@ -126,7 +128,7 @@ public class ShareGroupAssignorBenchmark {
setupTopics();
Map<String, ShareGroupMember> members = createMembers();
this.groupSpec = AssignorBenchmarkUtils.createShareGroupSpec(members, subscriptionType, topicResolver);
this.groupSpec = AssignorBenchmarkUtils.createShareGroupSpec(metadataImage, members, subscriptionType, topicResolver);
if (assignmentType == AssignmentType.INCREMENTAL) {
simulateIncrementalRebalance();
@@ -136,7 +138,7 @@ public class ShareGroupAssignorBenchmark {
private void setupTopics() {
allTopicNames = AssignorBenchmarkUtils.createTopicNames(topicCount);
CoordinatorMetadataImage metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionCount);
metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionCount);
topicResolver = new TopicIds.CachedTopicResolver(metadataImage);
subscribedTopicDescriber = new SubscribedTopicDescriberImpl(metadataImage);
@@ -171,8 +173,6 @@ public class ShareGroupAssignorBenchmark {
GroupAssignment initialAssignment = partitionAssignor.assign(groupSpec, subscribedTopicDescriber);
Map<String, MemberAssignment> members = initialAssignment.members();
Map<Uuid, Map<Integer, String>> invertedTargetAssignment = AssignorBenchmarkUtils.computeInvertedTargetAssignment(initialAssignment);
Map<String, MemberSubscriptionAndAssignmentImpl> updatedMemberSpec = new HashMap<>();
for (String memberId : groupSpec.memberIds()) {
@@ -204,9 +204,9 @@ public class ShareGroupAssignorBenchmark {
));
groupSpec = new GroupSpecImpl(
metadataImage,
updatedMemberSpec,
subscriptionType,
invertedTargetAssignment
subscriptionType
);
}
@@ -82,6 +82,9 @@ public class TargetAssignmentBuilderBenchmark {
@Param({"HOMOGENEOUS", "HETEROGENEOUS"})
private SubscriptionType subscriptionType;
@Param({"false", "true"})
private boolean assignorOffload;
private static final Logger LOG = new LogContext().logger(TargetAssignmentBuilderBenchmark.class);
private static final String GROUP_ID = "benchmark-group";
@@ -95,12 +98,8 @@ public class TargetAssignmentBuilderBenchmark {
/** The number of homogeneous subgroups to create for the heterogeneous subscription case. */
private static final int MAX_BUCKET_COUNT = 5;
private GroupSpec groupSpec;
private TargetAssignmentBuilder.TargetAssignmentResult targetAssignmentResult;
private Map<Uuid, Map<Integer, String>> invertedTargetAssignment;
private List<String> allTopicNames = List.of();
private CoordinatorMetadataImage metadataImage;
@@ -122,7 +121,6 @@ public class TargetAssignmentBuilderBenchmark {
ConsumerGroup group = new ConsumerGroup(logContext, snapshotRegistry, GROUP_ID);
createMembers().values().forEach(group::updateMember);
generateMockInitialTargetAssignment(group);
invertedTargetAssignment = group.invertedTargetAssignment();
ConsumerGroupMember newMember = new ConsumerGroupMember.Builder("newMember")
.setSubscribedTopicNames(allTopicNames)
@@ -131,8 +129,7 @@ public class TargetAssignmentBuilderBenchmark {
updatedMembersAndTargetAssignment = new UpdatedMembersAndTargetAssignmentView<>(group.members(), Map.of(), group.targetAssignment(), ConsumerGroupMember::instanceId);
updatedMembersAndTargetAssignment.addOrUpdateMember(newMember.memberId(), newMember);
groupSpec = buildGroupSpecWithoutOffload();
targetAssignmentResult = buildAssignment();
targetAssignmentResult = buildGroupSpecAndAssignment();
}
private void setupTopics() {
@@ -148,6 +145,7 @@ public class TargetAssignmentBuilderBenchmark {
private void generateMockInitialTargetAssignment(ConsumerGroup group) {
GroupSpec initialGroupSpec = AssignorBenchmarkUtils.createConsumerGroupSpec(
metadataImage,
group.members(),
subscriptionType,
topicResolver
@@ -196,31 +194,18 @@ public class TargetAssignmentBuilderBenchmark {
@Threads(1)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
public List<CoordinatorRecord> build() {
groupSpec = buildGroupSpecWithoutOffload();
targetAssignmentResult = buildAssignment();
targetAssignmentResult = buildGroupSpecAndAssignment();
return buildAssignmentRecords();
}
@Benchmark
@Threads(1)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
public GroupSpec buildGroupSpecWithoutOffload() {
return buildGroupSpec(false);
}
@Benchmark
@Threads(1)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
public GroupSpec buildGroupSpecWithOffload() {
return buildGroupSpec(true);
}
private GroupSpec buildGroupSpec(boolean assignorOffload) {
public GroupSpec buildGroupSpec() {
return new GroupSpecBuilder.ConsumerGroupSpecBuilder()
.withMembers(updatedMembersAndTargetAssignment.members())
.withSubscriptionType(subscriptionType)
.withTargetAssignment(updatedMembersAndTargetAssignment.targetAssignment())
.withInvertedTargetAssignment(invertedTargetAssignment)
.withMetadataImage(metadataImage)
.withAssignorOffload(assignorOffload)
.build();
@@ -229,7 +214,11 @@ public class TargetAssignmentBuilderBenchmark {
@Benchmark
@Threads(1)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
public TargetAssignmentBuilder.TargetAssignmentResult buildAssignment() {
public TargetAssignmentBuilder.TargetAssignmentResult buildGroupSpecAndAssignment() {
// We have to build a new group spec every time. The group spec caches its assigned
// partitions on the first call to isPartitionAssigned, so reusing one would leave that
// cost out of the benchmark.
GroupSpec groupSpec = buildGroupSpec();
return new TargetAssignmentBuilder(GROUP_EPOCH, partitionAssignor)
.withTime(Time.SYSTEM)
.withMetadataImage(metadataImage)
@@ -241,6 +230,8 @@ public class TargetAssignmentBuilderBenchmark {
@Threads(1)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
public List<CoordinatorRecord> buildAssignmentRecords() {
// Assignor offload does not affect this benchmark, since it only affects building the group
// spec.
return new TargetAssignmentRecordsBuilder.ConsumerTargetAssignmentRecordsBuilder(LOG, GROUP_ID)
.withTargetAssignmentMetadata(targetAssignmentResult.targetAssignmentMetadata())
.withCurrentMemberIds(updatedMembersAndTargetAssignment.members().keySet())