From 96b9c9867edf357a1c011d06575b588457033f0a Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Thu, 1 May 2025 20:46:54 +0800 Subject: [PATCH 1/8] KAFKA-17747: Get rid of TopicMetadata in SubscribedTopicDescriberImpl Signed-off-by: PoAn Yang --- .../group/GroupMetadataManager.java | 8 +- .../modern/SubscribedTopicDescriberImpl.java | 87 ++++-- .../group/modern/TargetAssignmentBuilder.java | 43 ++- ...OptimizedUniformAssignmentBuilderTest.java | 195 ++++++------- .../group/assignor/RangeAssignorTest.java | 256 +++++++----------- .../group/assignor/SimpleAssignorTest.java | 208 ++++++-------- ...ormHeterogeneousAssignmentBuilderTest.java | 246 +++++++---------- .../modern/SubscribedTopicMetadataTest.java | 87 ++++-- .../modern/TargetAssignmentBuilderTest.java | 30 +- .../jmh/assignor/AssignorBenchmarkUtils.java | 66 ++--- .../assignor/ServerSideAssignorBenchmark.java | 19 +- .../TargetAssignmentBuilderBenchmark.java | 23 +- 12 files changed, 568 insertions(+), 700 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java index fb2c7c4c6af45..e5334d21bb726 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java @@ -3647,11 +3647,11 @@ private Assignment updateTargetAssignment( new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(group.groupId(), groupEpoch, consumerGroupAssignors.get(preferredServerAssignor)) .withMembers(group.members()) .withStaticMembers(group.staticMembers()) - .withSubscriptionMetadata(subscriptionMetadata) + .withSubscriptionTopicIdSet(subscriptionMetadata.values().stream().map(TopicMetadata::id).collect(Collectors.toSet())) .withSubscriptionType(subscriptionType) .withTargetAssignment(group.targetAssignment()) .withInvertedTargetAssignment(group.invertedTargetAssignment()) - .withTopicsImage(metadataImage.topics()) + .withMetadataImage(metadataImage) .withResolvedRegularExpressions(group.resolvedRegularExpressions()) .addOrUpdateMember(updatedMember.memberId(), updatedMember); @@ -3718,12 +3718,12 @@ private Assignment updateTargetAssignment( TargetAssignmentBuilder.ShareTargetAssignmentBuilder assignmentResultBuilder = new TargetAssignmentBuilder.ShareTargetAssignmentBuilder(group.groupId(), groupEpoch, shareGroupAssignor) .withMembers(group.members()) - .withSubscriptionMetadata(subscriptionMetadata) + .withSubscriptionTopicIdSet(subscriptionMetadata.values().stream().map(TopicMetadata::id).collect(Collectors.toSet())) .withSubscriptionType(subscriptionType) .withTargetAssignment(group.targetAssignment()) .withTopicAssignablePartitionsMap(initializedTopicPartitions) .withInvertedTargetAssignment(group.invertedTargetAssignment()) - .withTopicsImage(metadataImage.topics()) + .withMetadataImage(metadataImage) .addOrUpdateMember(updatedMember.memberId(), updatedMember); long startTimeMs = time.milliseconds(); diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java index 5a5f2336d14ff..5b6488b520018 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java @@ -19,12 +19,15 @@ import org.apache.kafka.common.Uuid; import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor; import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber; +import org.apache.kafka.image.MetadataImage; +import org.apache.kafka.image.TopicImage; +import org.apache.kafka.metadata.PartitionRegistration; +import java.util.HashSet; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; -import java.util.stream.Collectors; -import java.util.stream.IntStream; /** * The subscribed topic metadata class is used by the {@link PartitionAssignor} to obtain @@ -32,28 +35,42 @@ */ public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber { /** - * The topic Ids mapped to their corresponding {@link TopicMetadata} - * object, which contains topic and partition metadata. + * The topic Ids set. */ - private final Map topicMetadata; + private final Set subscriptionTopicIdSet; private final Map> topicPartitionAllowedMap; - public SubscribedTopicDescriberImpl(Map topicMetadata) { - this(topicMetadata, null); + /** + * The metadata image that contains the latest metadata information. + */ + private final MetadataImage metadataImage; + + public SubscribedTopicDescriberImpl(Set subscriptionTopicIdSet, MetadataImage metadataImage) { + this(subscriptionTopicIdSet, metadataImage, null); } - public SubscribedTopicDescriberImpl(Map topicMetadata, Map> topicPartitionAllowedMap) { - this.topicMetadata = Objects.requireNonNull(topicMetadata); + public SubscribedTopicDescriberImpl(Set subscriptionTopicIdSet, MetadataImage metadataImage, Map> topicPartitionAllowedMap) { + this.subscriptionTopicIdSet = Objects.requireNonNull(subscriptionTopicIdSet); + this.metadataImage = Objects.requireNonNull(metadataImage); this.topicPartitionAllowedMap = topicPartitionAllowedMap; } /** - * Map of topic Ids to topic metadata. + * Set of topic Ids. + * + * @return The set of topic Ids. + */ + public Set subscriptionTopicIdSet() { + return this.subscriptionTopicIdSet; + } + + /** + * The metadata image that contains the latest metadata information. * - * @return The map of topic Ids to topic metadata. + * @return The metadata image that contains the latest metadata information. */ - public Map topicMetadata() { - return this.topicMetadata; + public MetadataImage metadataImage() { + return this.metadataImage; } /** @@ -65,8 +82,11 @@ public Map topicMetadata() { */ @Override public int numPartitions(Uuid topicId) { - TopicMetadata topic = this.topicMetadata.get(topicId); - return topic == null ? -1 : topic.numPartitions(); + if (!subscriptionTopicIdSet.contains(topicId)) { + return -1; + } + TopicImage topicImage = this.metadataImage.topics().getTopic(topicId); + return topicImage == null ? -1 : topicImage.partitions().size(); } /** @@ -79,13 +99,30 @@ public int numPartitions(Uuid topicId) { */ @Override public Set racksForPartition(Uuid topicId, int partition) { + if (!subscriptionTopicIdSet.contains(topicId)) { + return Set.of(); + } + + TopicImage topic = metadataImage.topics().getTopic(topicId); + if (topic != null) { + PartitionRegistration partitionRegistration = topic.partitions().get(partition); + if (partitionRegistration != null) { + Set racks = new HashSet<>(); + for (int replica : partitionRegistration.replicas) { + Optional rackOptional = metadataImage.cluster().broker(replica).rack(); + // Only add the rack if it is available for the broker/replica. + rackOptional.ifPresent(racks::add); + } + return racks; + } + } return Set.of(); } /** * Returns a set of assignable partitions from the topic metadata. * If the allowed partition map is null, all the partitions in the corresponding - * topic metadata are returned for the argument topic id. If allowed map is empty, + * topic image are returned for the argument topic id. If allowed map is empty, * empty set is returned. * * @param topicId The uuid of the topic @@ -93,13 +130,17 @@ public Set racksForPartition(Uuid topicId, int partition) { */ @Override public Set assignablePartitions(Uuid topicId) { - TopicMetadata topic = this.topicMetadata.get(topicId); + if (!subscriptionTopicIdSet.contains(topicId)) { + return Set.of(); + } + + TopicImage topic = metadataImage.topics().getTopic(topicId); if (topic == null) { return Set.of(); } if (topicPartitionAllowedMap == null) { - return IntStream.range(0, topic.numPartitions()).boxed().collect(Collectors.toUnmodifiableSet()); + return Set.copyOf(topic.partitions().keySet()); } return topicPartitionAllowedMap.getOrDefault(topicId, Set.of()); @@ -110,18 +151,22 @@ public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; SubscribedTopicDescriberImpl that = (SubscribedTopicDescriberImpl) o; - return topicMetadata.equals(that.topicMetadata); + if (!subscriptionTopicIdSet.equals(that.subscriptionTopicIdSet)) return false; + return metadataImage.equals(that.metadataImage); } @Override public int hashCode() { - return topicMetadata.hashCode(); + int result = subscriptionTopicIdSet.hashCode(); + result = 31 * result + metadataImage.hashCode(); + return result; } @Override public String toString() { return "SubscribedTopicMetadata(" + - "topicMetadata=" + topicMetadata + + "subscriptionTopicIdSet=" + subscriptionTopicIdSet + + ", metadataImage=" + metadataImage + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java index 63197eaba5ad5..31263b02fd80b 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java @@ -27,7 +27,7 @@ import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember; import org.apache.kafka.coordinator.group.modern.consumer.ResolvedRegularExpression; import org.apache.kafka.coordinator.group.modern.share.ShareGroupMember; -import org.apache.kafka.image.TopicsImage; +import org.apache.kafka.image.MetadataImage; import java.util.ArrayList; import java.util.Collections; @@ -251,9 +251,9 @@ protected MemberSubscriptionAndAssignmentImpl newMemberSubscriptionAndAssignment private Map members = Map.of(); /** - * The subscription metadata. + * The subscription topic id set. */ - private Map subscriptionMetadata = Map.of(); + private Set subscriptionTopicIdSet = Set.of(); /** * The subscription type of the consumer group. @@ -272,9 +272,9 @@ protected MemberSubscriptionAndAssignmentImpl newMemberSubscriptionAndAssignment private Map> invertedTargetAssignment = Map.of(); /** - * The topics image. + * The metadata image. */ - private TopicsImage topicsImage = TopicsImage.EMPTY; + private MetadataImage metadataImage = MetadataImage.EMPTY; /** * The members which have been updated or deleted. Deleted members @@ -336,15 +336,15 @@ public U withStaticMembers( } /** - * Adds the subscription metadata to use. + * Adds the set of topic id to use. * - * @param subscriptionMetadata The subscription metadata. + * @param subscriptionTopicIdSet The set of topic id. * @return This object. */ - public U withSubscriptionMetadata( - Map subscriptionMetadata + public U withSubscriptionTopicIdSet( + Set subscriptionTopicIdSet ) { - this.subscriptionMetadata = subscriptionMetadata; + this.subscriptionTopicIdSet = subscriptionTopicIdSet; return self(); } @@ -388,15 +388,15 @@ public U withInvertedTargetAssignment( } /** - * Adds the topics image. + * Adds the metadata image. * - * @param topicsImage The topics image. + * @param metadataImage The metadata image. * @return This object. */ - public U withTopicsImage( - TopicsImage topicsImage + public U withMetadataImage( + MetadataImage metadataImage ) { - this.topicsImage = topicsImage; + this.metadataImage = metadataImage; return self(); } @@ -445,7 +445,7 @@ public U removeMember( */ public TargetAssignmentResult build() throws PartitionAssignorException { Map memberSpecs = new HashMap<>(); - TopicIds.TopicResolver topicResolver = new TopicIds.CachedTopicResolver(topicsImage); + TopicIds.TopicResolver topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics()); // Prepare the member spec for all members. members.forEach((memberId, member) -> @@ -479,15 +479,6 @@ public TargetAssignmentResult build() throws PartitionAssignorException { } }); - // Prepare the topic metadata. - Map topicMetadataMap = new HashMap<>(); - subscriptionMetadata.forEach((topicName, topicMetadata) -> - topicMetadataMap.put( - topicMetadata.id(), - topicMetadata - ) - ); - // Compute the assignment. GroupAssignment newGroupAssignment = assignor.assign( new GroupSpecImpl( @@ -495,7 +486,7 @@ public TargetAssignmentResult build() throws PartitionAssignorException { subscriptionType, invertedTargetAssignment ), - new SubscribedTopicDescriberImpl(topicMetadataMap, topicAssignablePartitionsMap) + new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage, topicAssignablePartitionsMap) ); // Compute delta from previous to new target assignment and create the relevant records. diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java index e47aadb482ad7..40a568282f2b5 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java @@ -17,6 +17,7 @@ package org.apache.kafka.coordinator.group.assignor; import org.apache.kafka.common.Uuid; +import org.apache.kafka.coordinator.group.MetadataImageBuilder; import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment; import org.apache.kafka.coordinator.group.api.assignor.GroupSpec; import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignorException; @@ -24,7 +25,7 @@ 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.coordinator.group.modern.TopicMetadata; +import org.apache.kafka.image.MetadataImage; import org.junit.jupiter.api.Test; @@ -36,6 +37,7 @@ import java.util.Optional; import java.util.Set; import java.util.TreeMap; +import java.util.stream.IntStream; import static org.apache.kafka.coordinator.group.AssignmentTestUtil.assertAssignment; import static org.apache.kafka.coordinator.group.AssignmentTestUtil.invertedTargetAssignment; @@ -61,15 +63,13 @@ public class OptimizedUniformAssignmentBuilderTest { @Test public void testOneMemberNoTopicSubscription() { + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addRacks() + .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of( - topic1Uuid, - new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - ) - ) + Set.of(topic1Uuid), + metadataImage ); Map members = Map.of( @@ -98,15 +98,13 @@ public void testOneMemberNoTopicSubscription() { @Test public void testOneMemberSubscribedToNonexistentTopic() { + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addRacks() + .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of( - topic1Uuid, - new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - ) - ) + Set.of(topic1Uuid), + metadataImage ); Map members = Map.of( @@ -131,17 +129,11 @@ public void testOneMemberSubscribedToNonexistentTopic() { @Test public void testFirstAssignmentTwoMembersTwoTopicsNoMemberRacks() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic3Uuid, topic3Name, 2) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -173,7 +165,10 @@ public void testFirstAssignmentTwoMembersTwoTopicsNoMemberRacks() { HOMOGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -186,12 +181,10 @@ public void testFirstAssignmentTwoMembersTwoTopicsNoMemberRacks() { @Test public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic3Uuid, topic3Name, 2) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -233,7 +226,10 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { HOMOGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -246,22 +242,21 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { @Test public void testValidityAndBalanceForLargeSampleSet() { - Map topicMetadata = new HashMap<>(); - for (int i = 1; i < 100; i++) { + MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder(); + Set topicIds = new HashSet<>(); + IntStream.range(1, 101).forEach(i -> { Uuid topicId = Uuid.randomUuid(); - topicMetadata.put(topicId, new TopicMetadata( - topicId, - "topic-" + i, - 3 - )); - } + metadataImageBuilder.addTopic(topicId, "topic-" + i, 3); + topicIds.add(topicId); + }); + MetadataImage metadataImage = metadataImageBuilder.addRacks().build(); Map members = new TreeMap<>(); for (int i = 1; i < 50; i++) { members.put("member" + i, new MemberSubscriptionAndAssignmentImpl( Optional.empty(), Optional.empty(), - topicMetadata.keySet(), + topicIds, Assignment.EMPTY )); } @@ -271,7 +266,10 @@ public void testValidityAndBalanceForLargeSampleSet() { HOMOGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + topicIds, + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -283,17 +281,11 @@ public void testValidityAndBalanceForLargeSampleSet() { @Test public void testReassignmentForTwoMembersTwoTopicsGivenUnbalancedPrevAssignment() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 3 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic2Uuid, topic2Name, 3) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -332,7 +324,10 @@ public void testReassignmentForTwoMembersTwoTopicsGivenUnbalancedPrevAssignment( HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -346,17 +341,11 @@ public void testReassignmentForTwoMembersTwoTopicsGivenUnbalancedPrevAssignment( @Test public void testReassignmentWhenPartitionsAreAddedForTwoMembersTwoTopics() { // Simulating adding partition to T1 and T2 - originally T1 -> 3 Partitions and T2 -> 3 Partitions - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 6 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 5 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 6) + .addTopic(topic2Uuid, topic2Name, 5) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -395,7 +384,10 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembersTwoTopics() { HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -408,17 +400,11 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembersTwoTopics() { @Test public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembersTwoTopics() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 3 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic2Uuid, topic2Name, 3) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -467,7 +453,10 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -480,17 +469,11 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe @Test public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeMembersTwoTopics() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 3 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic2Uuid, topic2Name, 3) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -531,7 +514,10 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -544,17 +530,11 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM @Test public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWithTwoMembersTwoTopics() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 2 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 2) + .addTopic(topic2Uuid, topic2Name, 2) + .addRacks() + .build(); // Initial subscriptions were [T1, T2] Map members = new TreeMap<>(); @@ -592,7 +572,10 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 5ab032e47ca9f..42bea530f8a15 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -17,6 +17,7 @@ package org.apache.kafka.coordinator.group.assignor; import org.apache.kafka.common.Uuid; +import org.apache.kafka.coordinator.group.MetadataImageBuilder; 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; @@ -28,7 +29,7 @@ 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.coordinator.group.modern.TopicMetadata; +import org.apache.kafka.image.MetadataImage; import org.junit.jupiter.api.Test; @@ -61,14 +62,8 @@ public class RangeAssignorTest { @Test public void testOneMemberNoTopic() { SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of( - topic1Uuid, - new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - ) - ) + Set.of(), + MetadataImage.EMPTY ); Map members = Map.of( @@ -102,15 +97,13 @@ public void testOneMemberNoTopic() { @Test public void testOneMemberSubscribedToNonExistentTopic() { + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addRacks() + .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of( - topic1Uuid, - new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - ) - ) + Set.of(topic1Uuid), + metadataImage ); Map members = Map.of( @@ -135,17 +128,11 @@ public void testOneMemberSubscribedToNonExistentTopic() { @Test public void testFirstAssignmentTwoMembersTwoTopicsSameSubscriptions() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic3Uuid, topic3Name, 2) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -168,7 +155,10 @@ public void testFirstAssignmentTwoMembersTwoTopicsSameSubscriptions() { HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -190,22 +180,12 @@ public void testFirstAssignmentTwoMembersTwoTopicsSameSubscriptions() { @Test public void testFirstAssignmentThreeMembersThreeTopicsDifferentSubscriptions() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 3 - )); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic2Uuid, topic2Name, 3) + .addTopic(topic3Uuid, topic3Name, 2) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -235,7 +215,10 @@ public void testFirstAssignmentThreeMembersThreeTopicsDifferentSubscriptions() { HETEROGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid, topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -260,17 +243,11 @@ public void testFirstAssignmentThreeMembersThreeTopicsDifferentSubscriptions() { @Test public void testFirstAssignmentNumMembersGreaterThanNumPartitions() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic3Uuid, topic3Name, 2) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -300,7 +277,10 @@ public void testFirstAssignmentNumMembersGreaterThanNumPartitions() { HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -326,15 +306,13 @@ public void testFirstAssignmentNumMembersGreaterThanNumPartitions() { @Test public void testStaticMembership() throws PartitionAssignorException { + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addRacks() + .build(); SubscribedTopicDescriber subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of( - topic1Uuid, - new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - ) - ) + Set.of(topic1Uuid), + metadataImage ); Map members = new TreeMap<>(); @@ -396,15 +374,13 @@ public void testStaticMembership() throws PartitionAssignorException { @Test public void testMixedStaticMembership() throws PartitionAssignorException { + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 5) + .addRacks() + .build(); SubscribedTopicDescriber subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of( - topic1Uuid, - new TopicMetadata( - topic1Uuid, - topic1Name, - 5 - ) - ) + Set.of(topic1Uuid), + metadataImage ); // Initialize members with instance Ids. @@ -480,17 +456,11 @@ public void testMixedStaticMembership() throws PartitionAssignorException { @Test public void testReassignmentNumMembersGreaterThanNumPartitionsWhenOneMemberAdded() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 2 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 2) + .addTopic(topic2Uuid, topic2Name, 2) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -527,7 +497,10 @@ public void testReassignmentNumMembersGreaterThanNumPartitionsWhenOneMemberAdded HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -552,17 +525,11 @@ public void testReassignmentNumMembersGreaterThanNumPartitionsWhenOneMemberAdded @Test public void testReassignmentWhenOnePartitionAddedForTwoMembersTwoTopics() { // Simulating adding a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 4 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 4 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 4) + .addTopic(topic2Uuid, topic2Name, 4) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -591,7 +558,10 @@ public void testReassignmentWhenOnePartitionAddedForTwoMembersTwoTopics() { HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -613,17 +583,11 @@ public void testReassignmentWhenOnePartitionAddedForTwoMembersTwoTopics() { @Test public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembersTwoTopics() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 3 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic2Uuid, topic2Name, 3) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -660,7 +624,10 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -687,17 +654,11 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe @Test public void testReassignmentWhenOneMemberAddedAndOnePartitionAfterInitialAssignmentWithTwoMembersTwoTopics() { // Add a new partition to topic 1, initially T1 -> 3 partitions - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 4 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 3 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 4) + .addTopic(topic2Uuid, topic2Name, 3) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -734,7 +695,10 @@ public void testReassignmentWhenOneMemberAddedAndOnePartitionAfterInitialAssignm HETEROGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -759,17 +723,11 @@ public void testReassignmentWhenOneMemberAddedAndOnePartitionAfterInitialAssignm @Test public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithTwoMembersTwoTopics() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 3 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic2Uuid, topic2Name, 3) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -790,7 +748,10 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithTwoMem HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -808,22 +769,12 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithTwoMem @Test public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignmentWithThreeMembersTwoTopics() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 3 - )); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic2Uuid, topic2Name, 3) + .addTopic(topic3Uuid, topic3Name, 2) + .addRacks() + .build(); // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3 // Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2 @@ -863,7 +814,10 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme HETEROGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid, topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java index 3f33868e99955..34244916953d0 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java @@ -17,6 +17,7 @@ package org.apache.kafka.coordinator.group.assignor; import org.apache.kafka.common.Uuid; +import org.apache.kafka.coordinator.group.MetadataImageBuilder; 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; @@ -26,7 +27,7 @@ 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.coordinator.group.modern.TopicMetadata; +import org.apache.kafka.image.MetadataImage; import org.apache.kafka.server.common.TopicIdPartition; import org.junit.jupiter.api.Test; @@ -74,7 +75,8 @@ public void testName() { @Test public void testAssignWithEmptyMembers() { SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of() + Set.of(), + MetadataImage.EMPTY ); GroupSpec groupSpec = new GroupSpecImpl( @@ -104,15 +106,13 @@ public void testAssignWithEmptyMembers() { @Test public void testAssignWithNoSubscribedTopic() { + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) + .addRacks() + .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of( - TOPIC_1_UUID, - new TopicMetadata( - TOPIC_1_UUID, - TOPIC_1_NAME, - 3 - ) - ) + Set.of(TOPIC_1_UUID), + metadataImage ); Map members = Map.of( @@ -141,15 +141,13 @@ public void testAssignWithNoSubscribedTopic() { @Test public void testAssignWithSubscribedToNonExistentTopic() { + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) + .addRacks() + .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of( - TOPIC_1_UUID, - new TopicMetadata( - TOPIC_1_UUID, - TOPIC_1_NAME, - 3 - ) - ) + Set.of(TOPIC_1_UUID), + metadataImage ); Map members = Map.of( @@ -174,17 +172,11 @@ public void testAssignWithSubscribedToNonExistentTopic() { @Test public void testAssignWithTwoMembersAndTwoTopicsHomogeneous() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(TOPIC_1_UUID, new TopicMetadata( - TOPIC_1_UUID, - TOPIC_1_NAME, - 3 - )); - topicMetadata.put(TOPIC_3_UUID, new TopicMetadata( - TOPIC_3_UUID, - TOPIC_3_NAME, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) + .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2) + .addRacks() + .build(); Map members = new HashMap<>(); @@ -211,7 +203,10 @@ public void testAssignWithTwoMembersAndTwoTopicsHomogeneous() { HOMOGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(TOPIC_1_UUID, TOPIC_3_UUID), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -239,23 +234,12 @@ public void testAssignWithTwoMembersAndTwoTopicsHomogeneous() { @Test public void testAssignWithThreeMembersThreeTopicsHeterogeneous() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(TOPIC_1_UUID, new TopicMetadata( - TOPIC_1_UUID, - TOPIC_1_NAME, - 3 - )); - - topicMetadata.put(TOPIC_2_UUID, new TopicMetadata( - TOPIC_2_UUID, - "topic2", - 3 - )); - topicMetadata.put(TOPIC_3_UUID, new TopicMetadata( - TOPIC_3_UUID, - TOPIC_3_NAME, - 2 - )); + MetadataImage 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) + .addRacks() + .build(); Set memberATopicsSubscription = new LinkedHashSet<>(); memberATopicsSubscription.add(TOPIC_1_UUID); @@ -291,7 +275,10 @@ public void testAssignWithThreeMembersThreeTopicsHeterogeneous() { HETEROGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(TOPIC_1_UUID, TOPIC_2_UUID, TOPIC_3_UUID), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -321,18 +308,11 @@ public void testAssignWithThreeMembersThreeTopicsHeterogeneous() { @Test public void testAssignWithOneMemberNoAssignedTopicHeterogeneous() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(TOPIC_1_UUID, new TopicMetadata( - TOPIC_1_UUID, - TOPIC_1_NAME, - 3 - )); - - topicMetadata.put(TOPIC_2_UUID, new TopicMetadata( - TOPIC_2_UUID, - "topic2", - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) + .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 2) + .addRacks() + .build(); Set memberATopicsSubscription = new LinkedHashSet<>(); memberATopicsSubscription.add(TOPIC_1_UUID); @@ -357,7 +337,10 @@ public void testAssignWithOneMemberNoAssignedTopicHeterogeneous() { HETEROGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(TOPIC_1_UUID, TOPIC_2_UUID), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -495,17 +478,11 @@ public void testRoundRobinAssignmentWithCountTooManyPartitions() { @Test public void testAssignWithCurrentAssignmentHomogeneous() { // Current assignment setup - Two members A, B subscribing to T1 and T2. - Map topicMetadata1 = new HashMap<>(); - topicMetadata1.put(TOPIC_1_UUID, new TopicMetadata( - TOPIC_1_UUID, - TOPIC_1_NAME, - 3 - )); - topicMetadata1.put(TOPIC_2_UUID, new TopicMetadata( - TOPIC_2_UUID, - TOPIC_2_NAME, - 2 - )); + MetadataImage metadataImage1 = new MetadataImageBuilder() + .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) + .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 2) + .addRacks() + .build(); Map members1 = new HashMap<>(); @@ -532,7 +509,10 @@ public void testAssignWithCurrentAssignmentHomogeneous() { HOMOGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl(topicMetadata1); + SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl( + Set.of(TOPIC_1_UUID, TOPIC_2_UUID), + metadataImage1 + ); GroupAssignment computedAssignment1 = assignor.assign( groupSpec1, @@ -558,17 +538,11 @@ public void testAssignWithCurrentAssignmentHomogeneous() { assertAssignment(expectedAssignment1, computedAssignment1); // New assignment setup - Three members A, B, C subscribing to T2 and T3. - Map topicMetadata2 = new HashMap<>(); - topicMetadata2.put(TOPIC_2_UUID, new TopicMetadata( - TOPIC_2_UUID, - TOPIC_2_NAME, - 2 - )); - topicMetadata2.put(TOPIC_3_UUID, new TopicMetadata( - TOPIC_3_UUID, - TOPIC_3_NAME, - 3 - )); + MetadataImage metadataImage2 = new MetadataImageBuilder() + .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 2) + .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 3) + .addRacks() + .build(); Map members2 = new HashMap<>(); @@ -607,7 +581,10 @@ public void testAssignWithCurrentAssignmentHomogeneous() { HOMOGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl(topicMetadata2); + SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl( + Set.of(TOPIC_2_UUID, TOPIC_3_UUID), + metadataImage2 + ); GroupAssignment computedAssignment2 = assignor.assign( groupSpec2, @@ -639,23 +616,12 @@ public void testAssignWithCurrentAssignmentHomogeneous() { @Test public void testAssignWithCurrentAssignmentHeterogeneous() { // Current assignment setup - 3 members A - {T1, T2}, B - {T3}, C - {T2, T3}. - Map topicMetadata1 = new HashMap<>(); - topicMetadata1.put(TOPIC_1_UUID, new TopicMetadata( - TOPIC_1_UUID, - TOPIC_1_NAME, - 3 - )); - - topicMetadata1.put(TOPIC_2_UUID, new TopicMetadata( - TOPIC_2_UUID, - TOPIC_2_NAME, - 3 - )); - topicMetadata1.put(TOPIC_3_UUID, new TopicMetadata( - TOPIC_3_UUID, - TOPIC_3_NAME, - 2 - )); + MetadataImage 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) + .addRacks() + .build(); Set memberATopicsSubscription1 = new LinkedHashSet<>(); memberATopicsSubscription1.add(TOPIC_1_UUID); @@ -691,7 +657,10 @@ public void testAssignWithCurrentAssignmentHeterogeneous() { HETEROGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl(topicMetadata1); + SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl( + Set.of(TOPIC_1_UUID, TOPIC_2_UUID, TOPIC_3_UUID), + metadataImage1 + ); GroupAssignment computedAssignment1 = assignor.assign( groupSpec1, @@ -720,27 +689,13 @@ public void testAssignWithCurrentAssignmentHeterogeneous() { // New assignment setup - 2 members A - {T1, T2, T3}, B - {T3, T4}. - Map topicMetadata2 = new HashMap<>(); - topicMetadata2.put(TOPIC_1_UUID, new TopicMetadata( - TOPIC_1_UUID, - TOPIC_1_NAME, - 3 - )); - topicMetadata2.put(TOPIC_2_UUID, new TopicMetadata( - TOPIC_2_UUID, - TOPIC_2_NAME, - 3 - )); - topicMetadata2.put(TOPIC_3_UUID, new TopicMetadata( - TOPIC_3_UUID, - TOPIC_3_NAME, - 2 - )); - topicMetadata2.put(TOPIC_4_UUID, new TopicMetadata( - TOPIC_4_UUID, - TOPIC_4_NAME, - 1 - )); + MetadataImage 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) + .addRacks() + .build(); Map members2 = new HashMap<>(); @@ -776,7 +731,10 @@ public void testAssignWithCurrentAssignmentHeterogeneous() { Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl(topicMetadata2); + SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl( + Set.of(TOPIC_1_UUID, TOPIC_2_UUID, TOPIC_3_UUID, TOPIC_4_UUID), + metadataImage2 + ); GroupAssignment computedAssignment2 = assignor.assign( groupSpec2, diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java index decffabb6f492..fd59feb5f7d18 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java @@ -17,6 +17,7 @@ package org.apache.kafka.coordinator.group.assignor; import org.apache.kafka.common.Uuid; +import org.apache.kafka.coordinator.group.MetadataImageBuilder; import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment; import org.apache.kafka.coordinator.group.api.assignor.GroupSpec; import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignorException; @@ -25,7 +26,7 @@ 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.coordinator.group.modern.TopicMetadata; +import org.apache.kafka.image.MetadataImage; import org.junit.jupiter.api.Test; @@ -89,15 +90,13 @@ public Collection memberIds() { @Test public void testTwoMembersNoTopicSubscription() { + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addRacks() + .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of( - topic1Uuid, - new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - ) - ) + Set.of(topic1Uuid), + metadataImage ); Map members = new TreeMap<>(); @@ -130,15 +129,13 @@ public void testTwoMembersNoTopicSubscription() { @Test public void testTwoMembersSubscribedToNonexistentTopics() { + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addRacks() + .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Map.of( - topic1Uuid, - new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - ) - ) + Set.of(topic1Uuid), + metadataImage ); Map members = new TreeMap<>(); @@ -168,17 +165,11 @@ public void testTwoMembersSubscribedToNonexistentTopics() { @Test public void testFirstAssignmentTwoMembersTwoTopics() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 6 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic3Uuid, topic3Name, 6) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -201,7 +192,10 @@ public void testFirstAssignmentTwoMembersTwoTopics() { HETEROGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -222,17 +216,11 @@ public void testFirstAssignmentTwoMembersTwoTopics() { @Test public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 1 - )); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 2) + .addTopic(topic3Uuid, topic3Name, 1) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -262,7 +250,10 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { HETEROGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -286,22 +277,12 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { @Test public void testReassignmentForTwoMembersThreeTopicsGivenUnbalancedPrevAssignment() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 6 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 4 - )); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 4 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 6) + .addTopic(topic2Uuid, topic2Name, 4) + .addTopic(topic3Uuid, topic3Name, 4) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -340,7 +321,10 @@ public void testReassignmentForTwoMembersThreeTopicsGivenUnbalancedPrevAssignmen HETEROGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid, topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -366,27 +350,15 @@ public void testReassignmentForTwoMembersThreeTopicsGivenUnbalancedPrevAssignmen @Test public void testReassignmentWhenPartitionsAreAddedForTwoMembers() { // Simulating adding partitions to T1, T2, T3 - originally T1 -> 4, T2 -> 3, T3 -> 2, T4 -> 3 - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 6 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 5 - )); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 3 - )); - topicMetadata.put(topic4Uuid, new TopicMetadata( - topic4Uuid, - topic4Name, - 3 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 6) + .addTopic(topic2Uuid, topic2Name, 5) + .addTopic(topic3Uuid, topic3Name, 3) + .addTopic(topic4Uuid, topic4Name, 3) + .addRacks() + .build(); + + Map members = new TreeMap<>(); @@ -415,7 +387,10 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembers() { HETEROGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid, topic3Uuid, topic4Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -438,17 +413,11 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembers() { @Test public void testReassignmentWhenOneMemberAddedAndPartitionsAddedTwoMembersTwoTopics() { // Initially T1 -> 3, T2 -> 3 partitions. - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 6 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 7 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 6) + .addTopic(topic2Uuid, topic2Name, 7) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -485,7 +454,10 @@ public void testReassignmentWhenOneMemberAddedAndPartitionsAddedTwoMembersTwoTop HETEROGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -509,22 +481,12 @@ public void testReassignmentWhenOneMemberAddedAndPartitionsAddedTwoMembersTwoTop @Test public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeMembersThreeTopics() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 8 - )); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 3 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic2Uuid, topic2Name, 8) + .addTopic(topic3Uuid, topic3Name, 3) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -554,7 +516,10 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM HETEROGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid, topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -575,17 +540,11 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM @Test public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWithTwoMembersTwoTopics() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 5 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addTopic(topic2Uuid, topic2Name, 5) + .addRacks() + .build(); // Initial subscriptions were [T1, T2] Map members = new TreeMap<>(); @@ -615,7 +574,10 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith HETEROGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -641,22 +603,12 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith */ @Test public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 2 - )); - topicMetadata.put(topic2Uuid, new TopicMetadata( - topic2Uuid, - topic2Name, - 2 - )); - topicMetadata.put(topic3Uuid, new TopicMetadata( - topic3Uuid, - topic3Name, - 2 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 2) + .addTopic(topic2Uuid, topic2Name, 2) + .addTopic(topic3Uuid, topic3Name, 2) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -683,7 +635,10 @@ public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions( HETEROGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid, topic2Uuid, topic3Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, @@ -705,12 +660,10 @@ public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions( @Test public void testFirstAssignmentWithTwoMembersIncludingOneWithoutSubscriptions() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 3 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 3) + .addRacks() + .build(); Map members = new TreeMap<>(); @@ -733,7 +686,10 @@ public void testFirstAssignmentWithTwoMembersIncludingOneWithoutSubscriptions() HETEROGENEOUS, Map.of() ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( + Set.of(topic1Uuid), + metadataImage + ); GroupAssignment computedAssignment = assignor.assign( groupSpec, diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java index c0fd4a5a400c8..cdd0bb2b5a4f9 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java @@ -17,13 +17,15 @@ package org.apache.kafka.coordinator.group.modern; import org.apache.kafka.common.Uuid; +import org.apache.kafka.coordinator.group.MetadataImageBuilder; +import org.apache.kafka.image.MetadataImage; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -import java.util.HashMap; import java.util.Map; import java.util.Set; +import java.util.stream.IntStream; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; @@ -31,79 +33,110 @@ public class SubscribedTopicMetadataTest { - private Map topicMetadataMap; + private Set subscriptionTopicIdSet; private SubscribedTopicDescriberImpl subscribedTopicMetadata; + private MetadataImage metadataImage; + private final int numPartitions = 5; @BeforeEach public void setUp() { - topicMetadataMap = new HashMap<>(); - for (int i = 0; i < 5; i++) { + MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder(); + IntStream.range(0, 5).forEach(i -> { Uuid topicId = Uuid.randomUuid(); String topicName = "topic" + i; - topicMetadataMap.put( - topicId, - new TopicMetadata(topicId, topicName, 5) - ); - } - subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadataMap); + metadataImageBuilder.addTopic(topicId, topicName, numPartitions); + }); + metadataImageBuilder.addRacks(); + metadataImage = metadataImageBuilder.addRacks().build(); + + subscriptionTopicIdSet = metadataImage.topics().topicsById().keySet(); + subscribedTopicMetadata = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage); } @Test public void testAttribute() { - assertEquals(topicMetadataMap, subscribedTopicMetadata.topicMetadata()); + assertEquals(subscriptionTopicIdSet, subscribedTopicMetadata.subscriptionTopicIdSet()); + assertEquals(metadataImage, subscribedTopicMetadata.metadataImage()); + } + + @Test + public void testSubscriptionTopicIdSetCannotBeNull() { + assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(null, metadataImage)); } @Test - public void testTopicMetadataCannotBeNull() { - assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(null)); + public void testMetadataImageCannotBeNull() { + assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, null)); } @Test public void testNumberOfPartitions() { Uuid topicId = Uuid.randomUuid(); - // Test -1 is returned when the topic Id doesn't exist. + // Test -1 is returned when the topic ID doesn't exist. assertEquals(-1, subscribedTopicMetadata.numPartitions(topicId)); - topicMetadataMap.put(topicId, new TopicMetadata(topicId, "topic6", 3)); + // Test that the correct number of partitions are returned for a given topic ID. + subscriptionTopicIdSet.forEach(id -> + // Test that the correct number of partitions are returned for a given topic ID. + assertEquals(numPartitions, subscribedTopicMetadata.numPartitions(id)) + ); + } + + @Test + public void testRacksForPartition() { + Uuid topicId = Uuid.randomUuid(); + + // Test empty set is returned when the topic ID doesn't exist. + assertEquals(Set.of(), subscribedTopicMetadata.racksForPartition(topicId, 0)); + subscriptionTopicIdSet.forEach(id -> { + // Test empty set is returned when the partition ID doesn't exist. + assertEquals(Set.of(), subscribedTopicMetadata.racksForPartition(id, 10)); - // Test that the correct number of partitions are returned for a given topic Id. - assertEquals(3, subscribedTopicMetadata.numPartitions(topicId)); + // Test that the correct racks of partition are returned for a given topic ID. + assertEquals(Set.of("rack0", "rack1"), subscribedTopicMetadata.racksForPartition(id, 0)); + }); } @Test public void testEquals() { - assertEquals(new SubscribedTopicDescriberImpl(topicMetadataMap), subscribedTopicMetadata); + assertEquals(new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage), subscribedTopicMetadata); - Map topicMetadataMap2 = new HashMap<>(); Uuid topicId = Uuid.randomUuid(); - topicMetadataMap2.put(topicId, new TopicMetadata(topicId, "newTopic", 5)); - assertNotEquals(new SubscribedTopicDescriberImpl(topicMetadataMap2), subscribedTopicMetadata); + MetadataImage metadataImage2 = new MetadataImageBuilder() + .addTopic(topicId, "newTopic", 5) + .addRacks() + .build(); + Set subscriptionTopicIdSet2 = Set.of(topicId); + assertNotEquals(new SubscribedTopicDescriberImpl(subscriptionTopicIdSet2, metadataImage2), subscribedTopicMetadata); } @Test public void testAssignablePartitions() { - // null allow map (all partitions assignable) - subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadataMap, null); String t1Name = "t1"; Uuid t1Id = Uuid.randomUuid(); - topicMetadataMap.put(t1Id, new TopicMetadata(t1Id, t1Name, 5)); + metadataImage = new MetadataImageBuilder().addTopic(t1Id, t1Name, numPartitions).addRacks().build(); + subscriptionTopicIdSet = Set.of(t1Id); + // null allow map (all partitions assignable) + subscribedTopicMetadata = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage, null); assertEquals(Set.of(0, 1, 2, 3, 4), subscribedTopicMetadata.assignablePartitions(t1Id)); // empty allow map (nothing assignable) - subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadataMap, Map.of()); + subscribedTopicMetadata = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage, Map.of()); assertEquals(Set.of(), subscribedTopicMetadata.assignablePartitions(t1Id)); // few assignable partitions subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - topicMetadataMap, + subscriptionTopicIdSet, + metadataImage, Map.of(t1Id, Set.of(0, 5)) ); assertEquals(Set.of(0, 5), subscribedTopicMetadata.assignablePartitions(t1Id)); // all assignable partitions subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - topicMetadataMap, + subscriptionTopicIdSet, + metadataImage, Map.of(t1Id, Set.of(0, 1, 2, 3, 4)) ); assertEquals(Set.of(0, 1, 2, 3, 4), subscribedTopicMetadata.assignablePartitions(t1Id)); diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java index f33eed569180a..b3f7b27c7a329 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java @@ -25,12 +25,13 @@ import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType; import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember; import org.apache.kafka.coordinator.group.modern.consumer.ResolvedRegularExpression; -import org.apache.kafka.image.TopicsImage; +import org.apache.kafka.image.MetadataImage; import org.junit.jupiter.api.Test; import java.util.Arrays; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; @@ -57,13 +58,13 @@ public static class TargetAssignmentBuilderTestContext { private final int groupEpoch; private final PartitionAssignor assignor = mock(PartitionAssignor.class); private final Map members = new HashMap<>(); - private final Map subscriptionMetadata = new HashMap<>(); + private final Set subscriptionTopicIdSet = new HashSet<>(); private final Map updatedMembers = new HashMap<>(); private final Map targetAssignment = new HashMap<>(); private final Map memberAssignments = new HashMap<>(); private final Map staticMembers = new HashMap<>(); private final Map resolvedRegularExpressions = new HashMap<>(); - private MetadataImageBuilder topicsImageBuilder = new MetadataImageBuilder(); + private MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder(); public TargetAssignmentBuilderTestContext( String groupId, @@ -123,12 +124,8 @@ public Uuid addTopicMetadata( int numPartitions ) { Uuid topicId = Uuid.randomUuid(); - subscriptionMetadata.put(topicName, new TopicMetadata( - topicId, - topicName, - numPartitions - )); - topicsImageBuilder = topicsImageBuilder.addTopic(topicId, topicName, numPartitions); + subscriptionTopicIdSet.add(topicId); + metadataImageBuilder = metadataImageBuilder.addTopic(topicId, topicName, numPartitions); return topicId; } @@ -218,8 +215,8 @@ private MemberSubscriptionAndAssignmentImpl newMemberSubscriptionAndAssignment( } public TargetAssignmentBuilder.TargetAssignmentResult build() { - TopicsImage topicsImage = topicsImageBuilder.build().topics(); - TopicIds.TopicResolver topicResolver = new TopicIds.CachedTopicResolver(topicsImage); + MetadataImage metadataImage = metadataImageBuilder.build(); + TopicIds.TopicResolver topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics()); // Prepare expected member specs. Map memberSubscriptions = new HashMap<>(); @@ -256,13 +253,8 @@ public TargetAssignmentBuilder.TargetAssignmentResult build() { } }); - // Prepare the expected topic metadata. - Map topicMetadataMap = new HashMap<>(); - subscriptionMetadata.forEach((topicName, topicMetadata) -> - topicMetadataMap.put(topicMetadata.id(), topicMetadata)); - // Prepare the expected subscription topic metadata. - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadataMap); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage); SubscriptionType subscriptionType = HOMOGENEOUS; // Prepare the member assignments per topic partition. @@ -286,11 +278,11 @@ public TargetAssignmentBuilder.TargetAssignmentResult build() { new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(groupId, groupEpoch, assignor) .withMembers(members) .withStaticMembers(staticMembers) - .withSubscriptionMetadata(subscriptionMetadata) + .withSubscriptionTopicIdSet(subscriptionTopicIdSet) .withSubscriptionType(subscriptionType) .withTargetAssignment(targetAssignment) .withInvertedTargetAssignment(invertedTargetAssignment) - .withTopicsImage(topicsImage) + .withMetadataImage(metadataImage) .withResolvedRegularExpressions(resolvedRegularExpressions); // Add the updated members or delete the deleted members. diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java index 9402728e1eba2..00709466516b2 100644 --- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java +++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java @@ -27,12 +27,10 @@ import org.apache.kafka.coordinator.group.modern.GroupSpecImpl; import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl; import org.apache.kafka.coordinator.group.modern.TopicIds; -import org.apache.kafka.coordinator.group.modern.TopicMetadata; import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember; import org.apache.kafka.image.MetadataDelta; import org.apache.kafka.image.MetadataImage; import org.apache.kafka.image.MetadataProvenance; -import org.apache.kafka.image.TopicsImage; import java.util.ArrayList; import java.util.Arrays; @@ -89,69 +87,39 @@ public static List createTopicNames(int topicCount) { } /** - * Creates a subscription metadata map for the given topics. + * Get the subscription topic ID set from the given metadata image. * - * @param topicNames The names of the topics. - * @param partitionsPerTopic The number of partitions per topic. - * @return The subscription metadata map. + * @param metadataImage The metadata image. + * @return The set of topic id. */ - public static Map createSubscriptionMetadata( - List topicNames, - int partitionsPerTopic - ) { - Map subscriptionMetadata = new HashMap<>(); - - for (String topicName : topicNames) { - Uuid topicId = Uuid.randomUuid(); - - TopicMetadata metadata = new TopicMetadata( - topicId, - topicName, - partitionsPerTopic - ); - subscriptionMetadata.put(topicName, metadata); - } - - return subscriptionMetadata; - } - - /** - * Creates a topic metadata map from the given subscription metadata. - * - * @param subscriptionMetadata The subscription metadata. - * @return The topic metadata map. - */ - public static Map createTopicMetadata( - Map subscriptionMetadata - ) { - Map topicMetadata = new HashMap<>((int) (subscriptionMetadata.size() / 0.75f + 1)); - for (Map.Entry entry : subscriptionMetadata.entrySet()) { - topicMetadata.put(entry.getValue().id(), entry.getValue()); - } - return topicMetadata; + public static Set subscriptionTopicIdSet(MetadataImage metadataImage) { + return metadataImage.topics().topicsById().keySet(); } /** - * Creates a TopicsImage from the given subscription metadata. + * Creates a TopicsImage from the given topic names. * - * @param subscriptionMetadata The subscription metadata. + * @param allTopicNames The topic names. + * @param partitionsPerTopic Number of partitions per topic. * @return A TopicsImage containing the topic ids, names and partition counts from the * subscription metadata. */ - public static TopicsImage createTopicsImage(Map subscriptionMetadata) { + public static MetadataImage createMetadataImage( + List allTopicNames, + int partitionsPerTopic + ) { MetadataDelta delta = new MetadataDelta(MetadataImage.EMPTY); - for (Map.Entry entry : subscriptionMetadata.entrySet()) { - TopicMetadata topicMetadata = entry.getValue(); + for (String topicName : allTopicNames) { AssignorBenchmarkUtils.addTopic( delta, - topicMetadata.id(), - topicMetadata.name(), - topicMetadata.numPartitions() + Uuid.randomUuid(), + topicName, + partitionsPerTopic ); } - return delta.apply(MetadataProvenance.EMPTY).topics(); + return delta.apply(MetadataProvenance.EMPTY); } /** diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java index 971828c1c4726..4ad6dd9edb584 100644 --- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java +++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java @@ -31,9 +31,8 @@ import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl; import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl; import org.apache.kafka.coordinator.group.modern.TopicIds; -import org.apache.kafka.coordinator.group.modern.TopicMetadata; import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember; -import org.apache.kafka.image.TopicsImage; +import org.apache.kafka.image.MetadataImage; import org.openjdk.jmh.annotations.Benchmark; import org.openjdk.jmh.annotations.BenchmarkMode; @@ -125,9 +124,7 @@ public enum AssignmentType { private List allTopicNames = Collections.emptyList(); - private Map subscriptionMetadata = Collections.emptyMap(); - - private TopicsImage topicsImage = TopicsImage.EMPTY; + private MetadataImage metadataImage = MetadataImage.EMPTY; private TopicIds.TopicResolver topicResolver; @@ -151,16 +148,12 @@ private void setupTopics() { allTopicNames = AssignorBenchmarkUtils.createTopicNames(topicCount); int partitionsPerTopic = (memberCount * partitionsToMemberRatio) / topicCount; - subscriptionMetadata = AssignorBenchmarkUtils.createSubscriptionMetadata( - allTopicNames, - partitionsPerTopic - ); - topicsImage = AssignorBenchmarkUtils.createTopicsImage(subscriptionMetadata); - topicResolver = new TopicIds.CachedTopicResolver(topicsImage); + metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionsPerTopic); + topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics()); - Map topicMetadata = AssignorBenchmarkUtils.createTopicMetadata(subscriptionMetadata); - subscribedTopicDescriber = new SubscribedTopicDescriberImpl(topicMetadata); + Set subscriptionTopicIdSet = AssignorBenchmarkUtils.subscriptionTopicIdSet(metadataImage); + subscribedTopicDescriber = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage); } private Map createMembers() { diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java index 6fbb7908622b0..f87f00a5f1716 100644 --- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java +++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java @@ -28,9 +28,8 @@ import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl; import org.apache.kafka.coordinator.group.modern.TargetAssignmentBuilder; import org.apache.kafka.coordinator.group.modern.TopicIds; -import org.apache.kafka.coordinator.group.modern.TopicMetadata; import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember; -import org.apache.kafka.image.TopicsImage; +import org.apache.kafka.image.MetadataImage; import org.openjdk.jmh.annotations.Benchmark; import org.openjdk.jmh.annotations.BenchmarkMode; @@ -93,9 +92,9 @@ public class TargetAssignmentBuilderBenchmark { private List allTopicNames = Collections.emptyList(); - private Map subscriptionMetadata = Collections.emptyMap(); + private Set subscriptionTopicIdSet; - private TopicsImage topicsImage; + private MetadataImage metadataImage; private TopicIds.TopicResolver topicResolver; @@ -118,11 +117,11 @@ public void setup() { targetAssignmentBuilder = new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(GROUP_ID, GROUP_EPOCH, partitionAssignor) .withMembers(members) - .withSubscriptionMetadata(subscriptionMetadata) + .withSubscriptionTopicIdSet(subscriptionTopicIdSet) .withSubscriptionType(subscriptionType) .withTargetAssignment(existingTargetAssignment) .withInvertedTargetAssignment(invertedTargetAssignment) - .withTopicsImage(topicsImage) + .withMetadataImage(metadataImage) .addOrUpdateMember(newMember.memberId(), newMember); } @@ -130,16 +129,12 @@ private void setupTopics() { allTopicNames = AssignorBenchmarkUtils.createTopicNames(topicCount); int partitionsPerTopic = (memberCount * partitionsToMemberRatio) / topicCount; - subscriptionMetadata = AssignorBenchmarkUtils.createSubscriptionMetadata( - allTopicNames, - partitionsPerTopic - ); - topicsImage = AssignorBenchmarkUtils.createTopicsImage(subscriptionMetadata); - topicResolver = new TopicIds.CachedTopicResolver(topicsImage); + metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionsPerTopic); + topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics()); - Map topicMetadata = AssignorBenchmarkUtils.createTopicMetadata(subscriptionMetadata); - subscribedTopicDescriber = new SubscribedTopicDescriberImpl(topicMetadata); + subscriptionTopicIdSet = AssignorBenchmarkUtils.subscriptionTopicIdSet(metadataImage); + subscribedTopicDescriber = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage); } private Map generateMockInitialTargetAssignmentAndUpdateInvertedTargetAssignment( From be2c016f1c5f7e4fafca8da2058625c9a545a6c3 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Tue, 13 May 2025 19:12:00 +0800 Subject: [PATCH 2/8] address comment Signed-off-by: PoAn Yang --- .../group/GroupMetadataManager.java | 38 +++++++++---------- .../modern/SubscribedTopicDescriberImpl.java | 18 --------- .../modern/SubscribedTopicMetadataTest.java | 6 --- 3 files changed, 19 insertions(+), 43 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java index dbe83ce4e00f8..bb771045498ac 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java @@ -277,16 +277,16 @@ public class GroupMetadataManager { private static class UpdateSubscriptionMetadataResult { private final int groupEpoch; - private final Map subscriptionMetadata; + private final Set subscriptionTopicIdSet; private final SubscriptionType subscriptionType; UpdateSubscriptionMetadataResult( int groupEpoch, - Map subscriptionMetadata, + Set subscriptionTopicIdSet, SubscriptionType subscriptionType ) { this.groupEpoch = groupEpoch; - this.subscriptionMetadata = Objects.requireNonNull(subscriptionMetadata); + this.subscriptionTopicIdSet = Objects.requireNonNull(subscriptionTopicIdSet); this.subscriptionType = Objects.requireNonNull(subscriptionType); } } @@ -2223,8 +2223,8 @@ private CoordinatorResult ); int groupEpoch = group.groupEpoch(); - Map subscriptionMetadata = group.subscriptionMetadata(); SubscriptionType subscriptionType = group.subscriptionType(); + Set subscriptionTopicIdSet = group.subscriptionMetadata().values().stream().map(TopicMetadata::id).collect(Collectors.toSet()); if (bumpGroupEpoch || group.hasMetadataExpired(currentTimeMs)) { // The subscription metadata is updated in two cases: @@ -2239,7 +2239,7 @@ private CoordinatorResult ); groupEpoch = result.groupEpoch; - subscriptionMetadata = result.subscriptionMetadata; + subscriptionTopicIdSet = result.subscriptionTopicIdSet; subscriptionType = result.subscriptionType; } @@ -2254,7 +2254,7 @@ private CoordinatorResult groupEpoch, member, updatedMember, - subscriptionMetadata, + subscriptionTopicIdSet, subscriptionType, records ); @@ -2365,7 +2365,7 @@ private CoordinatorResult classicGroupJoinToConsumerGro } int groupEpoch = group.groupEpoch(); - Map subscriptionMetadata = group.subscriptionMetadata(); + Set subscriptionTopicIdSet = Set.of(); SubscriptionType subscriptionType = group.subscriptionType(); final ConsumerProtocolSubscription subscription = deserializeSubscription(protocols); @@ -2408,7 +2408,7 @@ private CoordinatorResult classicGroupJoinToConsumerGro ); groupEpoch = result.groupEpoch; - subscriptionMetadata = result.subscriptionMetadata; + subscriptionTopicIdSet = result.subscriptionTopicIdSet; subscriptionType = result.subscriptionType; } @@ -2423,7 +2423,7 @@ private CoordinatorResult classicGroupJoinToConsumerGro groupEpoch, member, updatedMember, - subscriptionMetadata, + subscriptionTopicIdSet, subscriptionType, records ); @@ -3617,7 +3617,7 @@ private UpdateSubscriptionMetadataResult updateSubscriptionMetadata( return new UpdateSubscriptionMetadataResult( groupEpoch, - subscriptionMetadata, + subscriptionMetadata.values().stream().map(TopicMetadata::id).collect(Collectors.toSet()), subscriptionType ); } @@ -3625,13 +3625,13 @@ private UpdateSubscriptionMetadataResult updateSubscriptionMetadata( /** * Updates the target assignment according to the updated member and subscription metadata. * - * @param group The ConsumerGroup. - * @param groupEpoch The group epoch. - * @param member The existing member. - * @param updatedMember The updated member. - * @param subscriptionMetadata The subscription metadata. - * @param subscriptionType The group subscription type. - * @param records The list to accumulate any new records. + * @param group The ConsumerGroup. + * @param groupEpoch The group epoch. + * @param member The existing member. + * @param updatedMember The updated member. + * @param subscriptionTopicIdSet The subscription metadata. + * @param subscriptionType The group subscription type. + * @param records The list to accumulate any new records. * @return The new target assignment. */ private Assignment updateTargetAssignment( @@ -3639,7 +3639,7 @@ private Assignment updateTargetAssignment( int groupEpoch, ConsumerGroupMember member, ConsumerGroupMember updatedMember, - Map subscriptionMetadata, + Set subscriptionTopicIdSet, SubscriptionType subscriptionType, List records ) { @@ -3652,7 +3652,7 @@ private Assignment updateTargetAssignment( new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(group.groupId(), groupEpoch, consumerGroupAssignors.get(preferredServerAssignor)) .withMembers(group.members()) .withStaticMembers(group.staticMembers()) - .withSubscriptionTopicIdSet(subscriptionMetadata.values().stream().map(TopicMetadata::id).collect(Collectors.toSet())) + .withSubscriptionTopicIdSet(subscriptionTopicIdSet) .withSubscriptionType(subscriptionType) .withTargetAssignment(group.targetAssignment()) .withInvertedTargetAssignment(group.invertedTargetAssignment()) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java index 5b6488b520018..ad4dbbdb84f0a 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java @@ -55,24 +55,6 @@ public SubscribedTopicDescriberImpl(Set subscriptionTopicIdSet, MetadataIm this.topicPartitionAllowedMap = topicPartitionAllowedMap; } - /** - * Set of topic Ids. - * - * @return The set of topic Ids. - */ - public Set subscriptionTopicIdSet() { - return this.subscriptionTopicIdSet; - } - - /** - * The metadata image that contains the latest metadata information. - * - * @return The metadata image that contains the latest metadata information. - */ - public MetadataImage metadataImage() { - return this.metadataImage; - } - /** * The number of partitions for the given topic Id. * diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java index cdd0bb2b5a4f9..7b90b26c4e09c 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java @@ -53,12 +53,6 @@ public void setUp() { subscribedTopicMetadata = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage); } - @Test - public void testAttribute() { - assertEquals(subscriptionTopicIdSet, subscribedTopicMetadata.subscriptionTopicIdSet()); - assertEquals(metadataImage, subscribedTopicMetadata.metadataImage()); - } - @Test public void testSubscriptionTopicIdSetCannotBeNull() { assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(null, metadataImage)); From dd645fcf9496cab936f188c29f05348989e00b4e Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Thu, 15 May 2025 23:31:44 +0800 Subject: [PATCH 3/8] KAFKA-17747: remove subscriptionTopicIdSet Signed-off-by: PoAn Yang --- .../group/GroupMetadataManager.java | 39 ++++------- .../modern/SubscribedTopicDescriberImpl.java | 36 +++------- .../group/modern/TargetAssignmentBuilder.java | 20 +----- ...OptimizedUniformAssignmentBuilderTest.java | 22 ++----- .../group/assignor/RangeAssignorTest.java | 13 ---- .../group/assignor/SimpleAssignorTest.java | 66 +++++++++++++------ ...ormHeterogeneousAssignmentBuilderTest.java | 11 ---- .../modern/SubscribedTopicMetadataTest.java | 27 +++----- .../modern/TargetAssignmentBuilderTest.java | 6 +- .../jmh/assignor/AssignorBenchmarkUtils.java | 10 --- .../assignor/ServerSideAssignorBenchmark.java | 3 +- .../TargetAssignmentBuilderBenchmark.java | 6 +- 12 files changed, 84 insertions(+), 175 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java index c4674aa1974ea..8d2dd5f630807 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java @@ -277,16 +277,13 @@ public class GroupMetadataManager { private static class UpdateSubscriptionMetadataResult { private final int groupEpoch; - private final Set subscriptionTopicIdSet; private final SubscriptionType subscriptionType; UpdateSubscriptionMetadataResult( int groupEpoch, - Set subscriptionTopicIdSet, SubscriptionType subscriptionType ) { this.groupEpoch = groupEpoch; - this.subscriptionTopicIdSet = Objects.requireNonNull(subscriptionTopicIdSet); this.subscriptionType = Objects.requireNonNull(subscriptionType); } } @@ -2224,7 +2221,6 @@ private CoordinatorResult int groupEpoch = group.groupEpoch(); SubscriptionType subscriptionType = group.subscriptionType(); - Set subscriptionTopicIdSet = group.subscriptionMetadata().values().stream().map(TopicMetadata::id).collect(Collectors.toSet()); if (bumpGroupEpoch || group.hasMetadataExpired(currentTimeMs)) { // The subscription metadata is updated in two cases: @@ -2239,7 +2235,6 @@ private CoordinatorResult ); groupEpoch = result.groupEpoch; - subscriptionTopicIdSet = result.subscriptionTopicIdSet; subscriptionType = result.subscriptionType; } @@ -2254,7 +2249,6 @@ private CoordinatorResult groupEpoch, member, updatedMember, - subscriptionTopicIdSet, subscriptionType, records ); @@ -2365,7 +2359,6 @@ private CoordinatorResult classicGroupJoinToConsumerGro } int groupEpoch = group.groupEpoch(); - Set subscriptionTopicIdSet = Set.of(); SubscriptionType subscriptionType = group.subscriptionType(); final ConsumerProtocolSubscription subscription = deserializeSubscription(protocols); @@ -2408,7 +2401,6 @@ private CoordinatorResult classicGroupJoinToConsumerGro ); groupEpoch = result.groupEpoch; - subscriptionTopicIdSet = result.subscriptionTopicIdSet; subscriptionType = result.subscriptionType; } @@ -2423,7 +2415,6 @@ private CoordinatorResult classicGroupJoinToConsumerGro groupEpoch, member, updatedMember, - subscriptionTopicIdSet, subscriptionType, records ); @@ -2597,7 +2588,6 @@ private CoordinatorResult subscriptionTopicIdSet, SubscriptionType subscriptionType, List records ) { @@ -3651,7 +3638,6 @@ private Assignment updateTargetAssignment( new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(group.groupId(), groupEpoch, consumerGroupAssignors.get(preferredServerAssignor)) .withMembers(group.members()) .withStaticMembers(group.staticMembers()) - .withSubscriptionTopicIdSet(subscriptionTopicIdSet) .withSubscriptionType(subscriptionType) .withTargetAssignment(group.targetAssignment()) .withInvertedTargetAssignment(group.invertedTargetAssignment()) @@ -3698,19 +3684,17 @@ private Assignment updateTargetAssignment( /** * Updates the target assignment according to the updated member and subscription metadata. * - * @param group The ShareGroup. - * @param groupEpoch The group epoch. - * @param updatedMember The updated member. - * @param subscriptionMetadata The subscription metadata. - * @param subscriptionType The group subscription type. - * @param records The list to accumulate any new records. + * @param group The ShareGroup. + * @param groupEpoch The group epoch. + * @param updatedMember The updated member. + * @param subscriptionType The group subscription type. + * @param records The list to accumulate any new records. * @return The new target assignment. */ private Assignment updateTargetAssignment( ShareGroup group, int groupEpoch, ShareGroupMember updatedMember, - Map subscriptionMetadata, SubscriptionType subscriptionType, List records ) { @@ -3722,7 +3706,6 @@ private Assignment updateTargetAssignment( TargetAssignmentBuilder.ShareTargetAssignmentBuilder assignmentResultBuilder = new TargetAssignmentBuilder.ShareTargetAssignmentBuilder(group.groupId(), groupEpoch, shareGroupAssignor) .withMembers(group.members()) - .withSubscriptionTopicIdSet(subscriptionMetadata.values().stream().map(TopicMetadata::id).collect(Collectors.toSet())) .withSubscriptionType(subscriptionType) .withTargetAssignment(group.targetAssignment()) .withTopicAssignablePartitionsMap(initializedTopicPartitions) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java index ad4dbbdb84f0a..3465275b2e3ff 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java @@ -26,7 +26,6 @@ import java.util.HashSet; import java.util.Map; import java.util.Objects; -import java.util.Optional; import java.util.Set; /** @@ -34,10 +33,6 @@ * topic and partition metadata for the topics that the modern group is subscribed to. */ public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber { - /** - * The topic Ids set. - */ - private final Set subscriptionTopicIdSet; private final Map> topicPartitionAllowedMap; /** @@ -45,12 +40,11 @@ public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber { */ private final MetadataImage metadataImage; - public SubscribedTopicDescriberImpl(Set subscriptionTopicIdSet, MetadataImage metadataImage) { - this(subscriptionTopicIdSet, metadataImage, null); + public SubscribedTopicDescriberImpl(MetadataImage metadataImage) { + this(metadataImage, Map.of()); } - public SubscribedTopicDescriberImpl(Set subscriptionTopicIdSet, MetadataImage metadataImage, Map> topicPartitionAllowedMap) { - this.subscriptionTopicIdSet = Objects.requireNonNull(subscriptionTopicIdSet); + public SubscribedTopicDescriberImpl(MetadataImage metadataImage, Map> topicPartitionAllowedMap) { this.metadataImage = Objects.requireNonNull(metadataImage); this.topicPartitionAllowedMap = topicPartitionAllowedMap; } @@ -64,9 +58,6 @@ public SubscribedTopicDescriberImpl(Set subscriptionTopicIdSet, MetadataIm */ @Override public int numPartitions(Uuid topicId) { - if (!subscriptionTopicIdSet.contains(topicId)) { - return -1; - } TopicImage topicImage = this.metadataImage.topics().getTopic(topicId); return topicImage == null ? -1 : topicImage.partitions().size(); } @@ -81,19 +72,14 @@ public int numPartitions(Uuid topicId) { */ @Override public Set racksForPartition(Uuid topicId, int partition) { - if (!subscriptionTopicIdSet.contains(topicId)) { - return Set.of(); - } - TopicImage topic = metadataImage.topics().getTopic(topicId); if (topic != null) { PartitionRegistration partitionRegistration = topic.partitions().get(partition); if (partitionRegistration != null) { Set racks = new HashSet<>(); for (int replica : partitionRegistration.replicas) { - Optional rackOptional = metadataImage.cluster().broker(replica).rack(); // Only add the rack if it is available for the broker/replica. - rackOptional.ifPresent(racks::add); + metadataImage.cluster().broker(replica).rack().ifPresent(racks::add); } return racks; } @@ -112,10 +98,6 @@ public Set racksForPartition(Uuid topicId, int partition) { */ @Override public Set assignablePartitions(Uuid topicId) { - if (!subscriptionTopicIdSet.contains(topicId)) { - return Set.of(); - } - TopicImage topic = metadataImage.topics().getTopic(topicId); if (topic == null) { return Set.of(); @@ -133,22 +115,22 @@ public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; SubscribedTopicDescriberImpl that = (SubscribedTopicDescriberImpl) o; - if (!subscriptionTopicIdSet.equals(that.subscriptionTopicIdSet)) return false; + if (!topicPartitionAllowedMap.equals(that.topicPartitionAllowedMap)) return false; return metadataImage.equals(that.metadataImage); } @Override public int hashCode() { - int result = subscriptionTopicIdSet.hashCode(); - result = 31 * result + metadataImage.hashCode(); + int result = metadataImage.hashCode(); + result = 31 * result + topicPartitionAllowedMap.hashCode(); return result; } @Override public String toString() { return "SubscribedTopicMetadata(" + - "subscriptionTopicIdSet=" + subscriptionTopicIdSet + - ", metadataImage=" + metadataImage + + "metadataImage=" + metadataImage + + ", topicPartitionAllowedMap=" + topicPartitionAllowedMap + ')'; } } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java index 31263b02fd80b..9d9e2e71b487c 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java @@ -250,11 +250,6 @@ protected MemberSubscriptionAndAssignmentImpl newMemberSubscriptionAndAssignment */ private Map members = Map.of(); - /** - * The subscription topic id set. - */ - private Set subscriptionTopicIdSet = Set.of(); - /** * The subscription type of the consumer group. */ @@ -335,19 +330,6 @@ public U withStaticMembers( return self(); } - /** - * Adds the set of topic id to use. - * - * @param subscriptionTopicIdSet The set of topic id. - * @return This object. - */ - public U withSubscriptionTopicIdSet( - Set subscriptionTopicIdSet - ) { - this.subscriptionTopicIdSet = subscriptionTopicIdSet; - return self(); - } - /** * Adds the subscription type in use. * @@ -486,7 +468,7 @@ public TargetAssignmentResult build() throws PartitionAssignorException { subscriptionType, invertedTargetAssignment ), - new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage, topicAssignablePartitionsMap) + new SubscribedTopicDescriberImpl(metadataImage, topicAssignablePartitionsMap) ); // Compute delta from previous to new target assignment and create the relevant records. diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java index b53f42063fa5c..b7375f99d82fd 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java @@ -68,7 +68,6 @@ public void testOneMemberNoTopicSubscription() { .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid), metadataImage ); @@ -103,7 +102,6 @@ public void testOneMemberSubscribedToNonexistentTopic() { .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid), metadataImage ); @@ -165,7 +163,6 @@ public void testFirstAssignmentTwoMembersTwoTopicsNoMemberRacks() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic3Uuid), metadataImage ); @@ -226,7 +223,6 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic3Uuid), metadataImage ); @@ -266,7 +262,6 @@ public void testValidityAndBalanceForLargeSampleSet() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - topicIds, metadataImage ); @@ -324,7 +319,6 @@ public void testReassignmentForTwoMembersTwoTopicsGivenUnbalancedPrevAssignment( invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -384,7 +378,6 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembersTwoTopics() { invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -453,7 +446,6 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -514,7 +506,6 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -572,7 +563,6 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -587,12 +577,10 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith @Test public void testReassignmentStickinessWhenAlreadyBalanced() { - Map topicMetadata = new HashMap<>(); - topicMetadata.put(topic1Uuid, new TopicMetadata( - topic1Uuid, - topic1Name, - 5 - )); + MetadataImage metadataImage = new MetadataImageBuilder() + .addTopic(topic1Uuid, topic1Name, 5) + .addRacks() + .build(); // A TreeMap ensures that memberA is first in the iteration order. Map members = new TreeMap<>(); @@ -643,7 +631,7 @@ public void testReassignmentStickinessWhenAlreadyBalanced() { HOMOGENEOUS, invertedTargetAssignment(members) ); - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage); GroupAssignment computedAssignment = assignor.assign( groupSpec, diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 42bea530f8a15..66de197d2b40d 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -62,7 +62,6 @@ public class RangeAssignorTest { @Test public void testOneMemberNoTopic() { SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(), MetadataImage.EMPTY ); @@ -102,7 +101,6 @@ public void testOneMemberSubscribedToNonExistentTopic() { .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid), metadataImage ); @@ -156,7 +154,6 @@ public void testFirstAssignmentTwoMembersTwoTopicsSameSubscriptions() { invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic3Uuid), metadataImage ); @@ -216,7 +213,6 @@ public void testFirstAssignmentThreeMembersThreeTopicsDifferentSubscriptions() { invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid, topic3Uuid), metadataImage ); @@ -278,7 +274,6 @@ public void testFirstAssignmentNumMembersGreaterThanNumPartitions() { invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic3Uuid), metadataImage ); @@ -311,7 +306,6 @@ public void testStaticMembership() throws PartitionAssignorException { .addRacks() .build(); SubscribedTopicDescriber subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid), metadataImage ); @@ -379,7 +373,6 @@ public void testMixedStaticMembership() throws PartitionAssignorException { .addRacks() .build(); SubscribedTopicDescriber subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid), metadataImage ); @@ -498,7 +491,6 @@ public void testReassignmentNumMembersGreaterThanNumPartitionsWhenOneMemberAdded invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -559,7 +551,6 @@ public void testReassignmentWhenOnePartitionAddedForTwoMembersTwoTopics() { invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -625,7 +616,6 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -696,7 +686,6 @@ public void testReassignmentWhenOneMemberAddedAndOnePartitionAfterInitialAssignm invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -749,7 +738,6 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithTwoMem invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -815,7 +803,6 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid, topic3Uuid), metadataImage ); diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java index 34244916953d0..1b107b840c9f1 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java @@ -75,7 +75,6 @@ public void testName() { @Test public void testAssignWithEmptyMembers() { SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(), MetadataImage.EMPTY ); @@ -111,8 +110,10 @@ public void testAssignWithNoSubscribedTopic() { .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(TOPIC_1_UUID), - metadataImage + metadataImage, + Map.of( + TOPIC_1_UUID, Set.of(0, 1, 2) + ) ); Map members = Map.of( @@ -146,8 +147,10 @@ public void testAssignWithSubscribedToNonExistentTopic() { .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(TOPIC_1_UUID), - metadataImage + metadataImage, + Map.of( + TOPIC_1_UUID, Set.of(0, 1, 2) + ) ); Map members = Map.of( @@ -204,8 +207,11 @@ public void testAssignWithTwoMembersAndTwoTopicsHomogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(TOPIC_1_UUID, TOPIC_3_UUID), - metadataImage + metadataImage, + Map.of( + TOPIC_1_UUID, Set.of(0, 1, 2), + TOPIC_3_UUID, Set.of(0, 1) + ) ); GroupAssignment computedAssignment = assignor.assign( @@ -276,8 +282,12 @@ public void testAssignWithThreeMembersThreeTopicsHeterogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(TOPIC_1_UUID, TOPIC_2_UUID, TOPIC_3_UUID), - metadataImage + metadataImage, + Map.of( + TOPIC_1_UUID, Set.of(0, 1, 2), + TOPIC_2_UUID, Set.of(0, 1, 2), + TOPIC_3_UUID, Set.of(0, 1) + ) ); GroupAssignment computedAssignment = assignor.assign( @@ -338,8 +348,11 @@ public void testAssignWithOneMemberNoAssignedTopicHeterogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(TOPIC_1_UUID, TOPIC_2_UUID), - metadataImage + metadataImage, + Map.of( + TOPIC_1_UUID, Set.of(0, 1, 2), + TOPIC_2_UUID, Set.of(0, 1) + ) ); GroupAssignment computedAssignment = assignor.assign( @@ -510,8 +523,11 @@ public void testAssignWithCurrentAssignmentHomogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl( - Set.of(TOPIC_1_UUID, TOPIC_2_UUID), - metadataImage1 + metadataImage1, + Map.of( + TOPIC_1_UUID, Set.of(0, 1, 2), + TOPIC_2_UUID, Set.of(0, 1) + ) ); GroupAssignment computedAssignment1 = assignor.assign( @@ -582,8 +598,11 @@ public void testAssignWithCurrentAssignmentHomogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl( - Set.of(TOPIC_2_UUID, TOPIC_3_UUID), - metadataImage2 + metadataImage2, + Map.of( + TOPIC_2_UUID, Set.of(0, 1), + TOPIC_3_UUID, Set.of(0, 1, 2) + ) ); GroupAssignment computedAssignment2 = assignor.assign( @@ -658,8 +677,12 @@ public void testAssignWithCurrentAssignmentHeterogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl( - Set.of(TOPIC_1_UUID, TOPIC_2_UUID, TOPIC_3_UUID), - metadataImage1 + metadataImage1, + Map.of( + TOPIC_1_UUID, Set.of(0, 1, 2), + TOPIC_2_UUID, Set.of(0, 1, 2), + TOPIC_3_UUID, Set.of(0, 1) + ) ); GroupAssignment computedAssignment1 = assignor.assign( @@ -732,8 +755,13 @@ public void testAssignWithCurrentAssignmentHeterogeneous() { ); SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl( - Set.of(TOPIC_1_UUID, TOPIC_2_UUID, TOPIC_3_UUID, TOPIC_4_UUID), - metadataImage2 + metadataImage2, + Map.of( + TOPIC_1_UUID, Set.of(0, 1, 2), + TOPIC_2_UUID, Set.of(0, 1, 2), + TOPIC_3_UUID, Set.of(0, 1), + TOPIC_4_UUID, Set.of(0) + ) ); GroupAssignment computedAssignment2 = assignor.assign( diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java index fd59feb5f7d18..7e6be4a5a3392 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java @@ -95,7 +95,6 @@ public void testTwoMembersNoTopicSubscription() { .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid), metadataImage ); @@ -134,7 +133,6 @@ public void testTwoMembersSubscribedToNonexistentTopics() { .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid), metadataImage ); @@ -193,7 +191,6 @@ public void testFirstAssignmentTwoMembersTwoTopics() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic3Uuid), metadataImage ); @@ -251,7 +248,6 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic3Uuid), metadataImage ); @@ -322,7 +318,6 @@ public void testReassignmentForTwoMembersThreeTopicsGivenUnbalancedPrevAssignmen invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid, topic3Uuid), metadataImage ); @@ -388,7 +383,6 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembers() { invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid, topic3Uuid, topic4Uuid), metadataImage ); @@ -455,7 +449,6 @@ public void testReassignmentWhenOneMemberAddedAndPartitionsAddedTwoMembersTwoTop invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -517,7 +510,6 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid, topic3Uuid), metadataImage ); @@ -575,7 +567,6 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid), metadataImage ); @@ -636,7 +627,6 @@ public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions( invertedTargetAssignment(members) ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid, topic2Uuid, topic3Uuid), metadataImage ); @@ -687,7 +677,6 @@ public void testFirstAssignmentWithTwoMembersIncludingOneWithoutSubscriptions() Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - Set.of(topic1Uuid), metadataImage ); diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java index 7b90b26c4e09c..8453b3cb36b2e 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java @@ -33,7 +33,6 @@ public class SubscribedTopicMetadataTest { - private Set subscriptionTopicIdSet; private SubscribedTopicDescriberImpl subscribedTopicMetadata; private MetadataImage metadataImage; private final int numPartitions = 5; @@ -49,18 +48,12 @@ public void setUp() { metadataImageBuilder.addRacks(); metadataImage = metadataImageBuilder.addRacks().build(); - subscriptionTopicIdSet = metadataImage.topics().topicsById().keySet(); - subscribedTopicMetadata = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage); - } - - @Test - public void testSubscriptionTopicIdSetCannotBeNull() { - assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(null, metadataImage)); + subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage); } @Test public void testMetadataImageCannotBeNull() { - assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, null)); + assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(null)); } @Test @@ -71,7 +64,7 @@ public void testNumberOfPartitions() { assertEquals(-1, subscribedTopicMetadata.numPartitions(topicId)); // Test that the correct number of partitions are returned for a given topic ID. - subscriptionTopicIdSet.forEach(id -> + metadataImage.topics().topicsById().forEach((id, name) -> // Test that the correct number of partitions are returned for a given topic ID. assertEquals(numPartitions, subscribedTopicMetadata.numPartitions(id)) ); @@ -83,7 +76,7 @@ public void testRacksForPartition() { // Test empty set is returned when the topic ID doesn't exist. assertEquals(Set.of(), subscribedTopicMetadata.racksForPartition(topicId, 0)); - subscriptionTopicIdSet.forEach(id -> { + metadataImage.topics().topicsById().forEach((id, name) -> { // Test empty set is returned when the partition ID doesn't exist. assertEquals(Set.of(), subscribedTopicMetadata.racksForPartition(id, 10)); @@ -94,15 +87,14 @@ public void testRacksForPartition() { @Test public void testEquals() { - assertEquals(new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage), subscribedTopicMetadata); + assertEquals(new SubscribedTopicDescriberImpl(metadataImage), subscribedTopicMetadata); Uuid topicId = Uuid.randomUuid(); MetadataImage metadataImage2 = new MetadataImageBuilder() .addTopic(topicId, "newTopic", 5) .addRacks() .build(); - Set subscriptionTopicIdSet2 = Set.of(topicId); - assertNotEquals(new SubscribedTopicDescriberImpl(subscriptionTopicIdSet2, metadataImage2), subscribedTopicMetadata); + assertNotEquals(new SubscribedTopicDescriberImpl(metadataImage2), subscribedTopicMetadata); } @Test @@ -110,18 +102,16 @@ public void testAssignablePartitions() { String t1Name = "t1"; Uuid t1Id = Uuid.randomUuid(); metadataImage = new MetadataImageBuilder().addTopic(t1Id, t1Name, numPartitions).addRacks().build(); - subscriptionTopicIdSet = Set.of(t1Id); // null allow map (all partitions assignable) - subscribedTopicMetadata = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage, null); + subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage, null); assertEquals(Set.of(0, 1, 2, 3, 4), subscribedTopicMetadata.assignablePartitions(t1Id)); // empty allow map (nothing assignable) - subscribedTopicMetadata = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage, Map.of()); + subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage, Map.of()); assertEquals(Set.of(), subscribedTopicMetadata.assignablePartitions(t1Id)); // few assignable partitions subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - subscriptionTopicIdSet, metadataImage, Map.of(t1Id, Set.of(0, 5)) ); @@ -129,7 +119,6 @@ public void testAssignablePartitions() { // all assignable partitions subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - subscriptionTopicIdSet, metadataImage, Map.of(t1Id, Set.of(0, 1, 2, 3, 4)) ); diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java index b3f7b27c7a329..59541ebad6a23 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java @@ -31,7 +31,6 @@ import java.util.Arrays; import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; @@ -58,7 +57,6 @@ public static class TargetAssignmentBuilderTestContext { private final int groupEpoch; private final PartitionAssignor assignor = mock(PartitionAssignor.class); private final Map members = new HashMap<>(); - private final Set subscriptionTopicIdSet = new HashSet<>(); private final Map updatedMembers = new HashMap<>(); private final Map targetAssignment = new HashMap<>(); private final Map memberAssignments = new HashMap<>(); @@ -124,7 +122,6 @@ public Uuid addTopicMetadata( int numPartitions ) { Uuid topicId = Uuid.randomUuid(); - subscriptionTopicIdSet.add(topicId); metadataImageBuilder = metadataImageBuilder.addTopic(topicId, topicName, numPartitions); return topicId; @@ -254,7 +251,7 @@ public TargetAssignmentBuilder.TargetAssignmentResult build() { }); // Prepare the expected subscription topic metadata. - SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage); + SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage); SubscriptionType subscriptionType = HOMOGENEOUS; // Prepare the member assignments per topic partition. @@ -278,7 +275,6 @@ public TargetAssignmentBuilder.TargetAssignmentResult build() { new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(groupId, groupEpoch, assignor) .withMembers(members) .withStaticMembers(staticMembers) - .withSubscriptionTopicIdSet(subscriptionTopicIdSet) .withSubscriptionType(subscriptionType) .withTargetAssignment(targetAssignment) .withInvertedTargetAssignment(invertedTargetAssignment) diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java index 00709466516b2..d6515a5f56536 100644 --- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java +++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java @@ -86,16 +86,6 @@ public static List createTopicNames(int topicCount) { return topicNames; } - /** - * Get the subscription topic ID set from the given metadata image. - * - * @param metadataImage The metadata image. - * @return The set of topic id. - */ - public static Set subscriptionTopicIdSet(MetadataImage metadataImage) { - return metadataImage.topics().topicsById().keySet(); - } - /** * Creates a TopicsImage from the given topic names. * diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java index 4ad6dd9edb584..fa232c9427fc7 100644 --- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java +++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java @@ -152,8 +152,7 @@ private void setupTopics() { metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionsPerTopic); topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics()); - Set subscriptionTopicIdSet = AssignorBenchmarkUtils.subscriptionTopicIdSet(metadataImage); - subscribedTopicDescriber = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage); + subscribedTopicDescriber = new SubscribedTopicDescriberImpl(metadataImage); } private Map createMembers() { diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java index f87f00a5f1716..c150d94551a89 100644 --- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java +++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java @@ -92,8 +92,6 @@ public class TargetAssignmentBuilderBenchmark { private List allTopicNames = Collections.emptyList(); - private Set subscriptionTopicIdSet; - private MetadataImage metadataImage; private TopicIds.TopicResolver topicResolver; @@ -117,7 +115,6 @@ public void setup() { targetAssignmentBuilder = new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(GROUP_ID, GROUP_EPOCH, partitionAssignor) .withMembers(members) - .withSubscriptionTopicIdSet(subscriptionTopicIdSet) .withSubscriptionType(subscriptionType) .withTargetAssignment(existingTargetAssignment) .withInvertedTargetAssignment(invertedTargetAssignment) @@ -133,8 +130,7 @@ private void setupTopics() { metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionsPerTopic); topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics()); - subscriptionTopicIdSet = AssignorBenchmarkUtils.subscriptionTopicIdSet(metadataImage); - subscribedTopicDescriber = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage); + subscribedTopicDescriber = new SubscribedTopicDescriberImpl(metadataImage); } private Map generateMockInitialTargetAssignmentAndUpdateInvertedTargetAssignment( From 3934672560b7dc6dc2d90d0d2dd1c3175daf340f Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Fri, 16 May 2025 14:11:48 +0800 Subject: [PATCH 4/8] KAFKA-17747: change topicPartitionAllowedMap to Optional Signed-off-by: PoAn Yang --- .../modern/SubscribedTopicDescriberImpl.java | 22 +++++--- .../group/modern/TargetAssignmentBuilder.java | 4 +- .../group/assignor/SimpleAssignorTest.java | 56 +++---------------- .../modern/SubscribedTopicMetadataTest.java | 11 ++-- 4 files changed, 32 insertions(+), 61 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java index 3465275b2e3ff..484f4e33d5f3f 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java @@ -26,6 +26,7 @@ import java.util.HashSet; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; /** @@ -33,7 +34,11 @@ * topic and partition metadata for the topics that the modern group is subscribed to. */ public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber { - private final Map> topicPartitionAllowedMap; + /** + * The map of topic Ids to the set of allowed partitions for each topic. + * If this is empty, all partitions are allowed. + */ + private final Optional>> topicPartitionAllowedMap; /** * The metadata image that contains the latest metadata information. @@ -41,10 +46,13 @@ public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber { private final MetadataImage metadataImage; public SubscribedTopicDescriberImpl(MetadataImage metadataImage) { - this(metadataImage, Map.of()); + this(metadataImage, Optional.empty()); } - public SubscribedTopicDescriberImpl(MetadataImage metadataImage, Map> topicPartitionAllowedMap) { + public SubscribedTopicDescriberImpl( + MetadataImage metadataImage, + Optional>> topicPartitionAllowedMap + ) { this.metadataImage = Objects.requireNonNull(metadataImage); this.topicPartitionAllowedMap = topicPartitionAllowedMap; } @@ -88,8 +96,8 @@ public Set racksForPartition(Uuid topicId, int partition) { } /** - * Returns a set of assignable partitions from the topic metadata. - * If the allowed partition map is null, all the partitions in the corresponding + * Returns a set of assignable partitions from the metadata image. + * If the allowed partition map is Optional.empty(), all the partitions in the corresponding * topic image are returned for the argument topic id. If allowed map is empty, * empty set is returned. * @@ -103,11 +111,11 @@ public Set assignablePartitions(Uuid topicId) { return Set.of(); } - if (topicPartitionAllowedMap == null) { + if (topicPartitionAllowedMap.isEmpty()) { return Set.copyOf(topic.partitions().keySet()); } - return topicPartitionAllowedMap.getOrDefault(topicId, Set.of()); + return topicPartitionAllowedMap.get().getOrDefault(topicId, Set.of()); } @Override diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java index 9d9e2e71b487c..ed1755e483e18 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java @@ -285,7 +285,7 @@ protected MemberSubscriptionAndAssignmentImpl newMemberSubscriptionAndAssignment /** * Topic partition assignable map. */ - private Map> topicAssignablePartitionsMap = new HashMap<>(); + private Optional>> topicAssignablePartitionsMap = Optional.empty(); /** * Constructs the object. @@ -385,7 +385,7 @@ public U withMetadataImage( public U withTopicAssignablePartitionsMap( Map> topicAssignablePartitionsMap ) { - this.topicAssignablePartitionsMap = topicAssignablePartitionsMap; + this.topicAssignablePartitionsMap = Optional.of(topicAssignablePartitionsMap); return self(); } diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java index 1b107b840c9f1..39fe200cd568f 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java @@ -110,10 +110,7 @@ public void testAssignWithNoSubscribedTopic() { .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - metadataImage, - Map.of( - TOPIC_1_UUID, Set.of(0, 1, 2) - ) + metadataImage ); Map members = Map.of( @@ -147,10 +144,7 @@ public void testAssignWithSubscribedToNonExistentTopic() { .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - metadataImage, - Map.of( - TOPIC_1_UUID, Set.of(0, 1, 2) - ) + metadataImage ); Map members = Map.of( @@ -207,11 +201,7 @@ public void testAssignWithTwoMembersAndTwoTopicsHomogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - metadataImage, - Map.of( - TOPIC_1_UUID, Set.of(0, 1, 2), - TOPIC_3_UUID, Set.of(0, 1) - ) + metadataImage ); GroupAssignment computedAssignment = assignor.assign( @@ -282,12 +272,7 @@ public void testAssignWithThreeMembersThreeTopicsHeterogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - metadataImage, - Map.of( - TOPIC_1_UUID, Set.of(0, 1, 2), - TOPIC_2_UUID, Set.of(0, 1, 2), - TOPIC_3_UUID, Set.of(0, 1) - ) + metadataImage ); GroupAssignment computedAssignment = assignor.assign( @@ -348,11 +333,7 @@ public void testAssignWithOneMemberNoAssignedTopicHeterogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( - metadataImage, - Map.of( - TOPIC_1_UUID, Set.of(0, 1, 2), - TOPIC_2_UUID, Set.of(0, 1) - ) + metadataImage ); GroupAssignment computedAssignment = assignor.assign( @@ -523,11 +504,7 @@ public void testAssignWithCurrentAssignmentHomogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl( - metadataImage1, - Map.of( - TOPIC_1_UUID, Set.of(0, 1, 2), - TOPIC_2_UUID, Set.of(0, 1) - ) + metadataImage1 ); GroupAssignment computedAssignment1 = assignor.assign( @@ -598,11 +575,7 @@ public void testAssignWithCurrentAssignmentHomogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl( - metadataImage2, - Map.of( - TOPIC_2_UUID, Set.of(0, 1), - TOPIC_3_UUID, Set.of(0, 1, 2) - ) + metadataImage2 ); GroupAssignment computedAssignment2 = assignor.assign( @@ -677,12 +650,7 @@ public void testAssignWithCurrentAssignmentHeterogeneous() { Map.of() ); SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl( - metadataImage1, - Map.of( - TOPIC_1_UUID, Set.of(0, 1, 2), - TOPIC_2_UUID, Set.of(0, 1, 2), - TOPIC_3_UUID, Set.of(0, 1) - ) + metadataImage1 ); GroupAssignment computedAssignment1 = assignor.assign( @@ -755,13 +723,7 @@ public void testAssignWithCurrentAssignmentHeterogeneous() { ); SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl( - metadataImage2, - Map.of( - TOPIC_1_UUID, Set.of(0, 1, 2), - TOPIC_2_UUID, Set.of(0, 1, 2), - TOPIC_3_UUID, Set.of(0, 1), - TOPIC_4_UUID, Set.of(0) - ) + metadataImage2 ); GroupAssignment computedAssignment2 = assignor.assign( diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java index 8453b3cb36b2e..47a8e43ff5117 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java @@ -24,6 +24,7 @@ import org.junit.jupiter.api.Test; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.stream.IntStream; @@ -102,25 +103,25 @@ public void testAssignablePartitions() { String t1Name = "t1"; Uuid t1Id = Uuid.randomUuid(); metadataImage = new MetadataImageBuilder().addTopic(t1Id, t1Name, numPartitions).addRacks().build(); - // null allow map (all partitions assignable) - subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage, null); + // Optional.empty() allow map (all partitions assignable) + subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage, Optional.empty()); assertEquals(Set.of(0, 1, 2, 3, 4), subscribedTopicMetadata.assignablePartitions(t1Id)); // empty allow map (nothing assignable) - subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage, Map.of()); + subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage, Optional.of(Map.of())); assertEquals(Set.of(), subscribedTopicMetadata.assignablePartitions(t1Id)); // few assignable partitions subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage, - Map.of(t1Id, Set.of(0, 5)) + Optional.of(Map.of(t1Id, Set.of(0, 5))) ); assertEquals(Set.of(0, 5), subscribedTopicMetadata.assignablePartitions(t1Id)); // all assignable partitions subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage, - Map.of(t1Id, Set.of(0, 1, 2, 3, 4)) + Optional.of(Map.of(t1Id, Set.of(0, 1, 2, 3, 4))) ); assertEquals(Set.of(0, 1, 2, 3, 4), subscribedTopicMetadata.assignablePartitions(t1Id)); } From bacd97d3c7634535c627c4927912973682dc4902 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Sun, 18 May 2025 13:31:00 +0800 Subject: [PATCH 5/8] address comment Signed-off-by: PoAn Yang --- .../modern/SubscribedTopicDescriberImpl.java | 2 +- .../OptimizedUniformAssignmentBuilderTest.java | 17 +++-------------- .../group/assignor/RangeAssignorTest.java | 12 ------------ .../group/assignor/SimpleAssignorTest.java | 9 --------- ...iformHeterogeneousAssignmentBuilderTest.java | 11 ----------- .../modern/SubscribedTopicMetadataTest.java | 13 ++++++++----- 6 files changed, 12 insertions(+), 52 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java index 484f4e33d5f3f..740769baf4198 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java @@ -54,7 +54,7 @@ public SubscribedTopicDescriberImpl( Optional>> topicPartitionAllowedMap ) { this.metadataImage = Objects.requireNonNull(metadataImage); - this.topicPartitionAllowedMap = topicPartitionAllowedMap; + this.topicPartitionAllowedMap = Objects.requireNonNull(topicPartitionAllowedMap); } /** diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java index b7375f99d82fd..1038c3833d710 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java @@ -37,7 +37,6 @@ import java.util.Optional; import java.util.Set; import java.util.TreeMap; -import java.util.stream.IntStream; import static org.apache.kafka.coordinator.group.AssignmentTestUtil.assertAssignment; import static org.apache.kafka.coordinator.group.AssignmentTestUtil.invertedTargetAssignment; @@ -65,7 +64,6 @@ public class OptimizedUniformAssignmentBuilderTest { public void testOneMemberNoTopicSubscription() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) - .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage @@ -99,7 +97,6 @@ public void testOneMemberNoTopicSubscription() { public void testOneMemberSubscribedToNonexistentTopic() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) - .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage @@ -130,7 +127,6 @@ public void testFirstAssignmentTwoMembersTwoTopicsNoMemberRacks() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic3Uuid, topic3Name, 2) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -179,7 +175,6 @@ public void testFirstAssignmentTwoMembersTwoTopicsNoMemberRacks() { public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic3Uuid, topic3Name, 2) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -239,12 +234,12 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { public void testValidityAndBalanceForLargeSampleSet() { MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder(); Set topicIds = new HashSet<>(); - IntStream.range(1, 101).forEach(i -> { + for (int i = 1; i < 101; i++) { Uuid topicId = Uuid.randomUuid(); metadataImageBuilder.addTopic(topicId, "topic-" + i, 3); topicIds.add(topicId); - }); - MetadataImage metadataImage = metadataImageBuilder.addRacks().build(); + } + MetadataImage metadataImage = metadataImageBuilder.build(); Map members = new TreeMap<>(); for (int i = 1; i < 50; i++) { @@ -278,7 +273,6 @@ public void testReassignmentForTwoMembersTwoTopicsGivenUnbalancedPrevAssignment( MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic2Uuid, topic2Name, 3) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -337,7 +331,6 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembersTwoTopics() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 6) .addTopic(topic2Uuid, topic2Name, 5) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -395,7 +388,6 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic2Uuid, topic2Name, 3) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -463,7 +455,6 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic2Uuid, topic2Name, 3) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -523,7 +514,6 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 2) .addTopic(topic2Uuid, topic2Name, 2) - .addRacks() .build(); // Initial subscriptions were [T1, T2] @@ -579,7 +569,6 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith public void testReassignmentStickinessWhenAlreadyBalanced() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 5) - .addRacks() .build(); // A TreeMap ensures that memberA is first in the iteration order. diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java index 66de197d2b40d..f83655a3196c7 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java @@ -98,7 +98,6 @@ public void testOneMemberNoTopic() { public void testOneMemberSubscribedToNonExistentTopic() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) - .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage @@ -129,7 +128,6 @@ public void testFirstAssignmentTwoMembersTwoTopicsSameSubscriptions() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic3Uuid, topic3Name, 2) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -181,7 +179,6 @@ public void testFirstAssignmentThreeMembersThreeTopicsDifferentSubscriptions() { .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic2Uuid, topic2Name, 3) .addTopic(topic3Uuid, topic3Name, 2) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -242,7 +239,6 @@ public void testFirstAssignmentNumMembersGreaterThanNumPartitions() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic3Uuid, topic3Name, 2) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -303,7 +299,6 @@ public void testFirstAssignmentNumMembersGreaterThanNumPartitions() { public void testStaticMembership() throws PartitionAssignorException { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) - .addRacks() .build(); SubscribedTopicDescriber subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage @@ -370,7 +365,6 @@ public void testStaticMembership() throws PartitionAssignorException { public void testMixedStaticMembership() throws PartitionAssignorException { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 5) - .addRacks() .build(); SubscribedTopicDescriber subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage @@ -452,7 +446,6 @@ public void testReassignmentNumMembersGreaterThanNumPartitionsWhenOneMemberAdded MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 2) .addTopic(topic2Uuid, topic2Name, 2) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -520,7 +513,6 @@ public void testReassignmentWhenOnePartitionAddedForTwoMembersTwoTopics() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 4) .addTopic(topic2Uuid, topic2Name, 4) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -577,7 +569,6 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic2Uuid, topic2Name, 3) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -647,7 +638,6 @@ public void testReassignmentWhenOneMemberAddedAndOnePartitionAfterInitialAssignm MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 4) .addTopic(topic2Uuid, topic2Name, 3) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -715,7 +705,6 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithTwoMem MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic2Uuid, topic2Name, 3) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -761,7 +750,6 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic2Uuid, topic2Name, 3) .addTopic(topic3Uuid, topic3Name, 2) - .addRacks() .build(); // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3 diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java index 39fe200cd568f..01ad2f17d05b7 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java @@ -107,7 +107,6 @@ public void testAssignWithEmptyMembers() { public void testAssignWithNoSubscribedTopic() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) - .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage @@ -141,7 +140,6 @@ public void testAssignWithNoSubscribedTopic() { public void testAssignWithSubscribedToNonExistentTopic() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) - .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage @@ -172,7 +170,6 @@ public void testAssignWithTwoMembersAndTwoTopicsHomogeneous() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2) - .addRacks() .build(); Map members = new HashMap<>(); @@ -234,7 +231,6 @@ public void testAssignWithThreeMembersThreeTopicsHeterogeneous() { .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3) .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2) - .addRacks() .build(); Set memberATopicsSubscription = new LinkedHashSet<>(); @@ -306,7 +302,6 @@ public void testAssignWithOneMemberNoAssignedTopicHeterogeneous() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 2) - .addRacks() .build(); Set memberATopicsSubscription = new LinkedHashSet<>(); @@ -475,7 +470,6 @@ public void testAssignWithCurrentAssignmentHomogeneous() { MetadataImage metadataImage1 = new MetadataImageBuilder() .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 2) - .addRacks() .build(); Map members1 = new HashMap<>(); @@ -534,7 +528,6 @@ public void testAssignWithCurrentAssignmentHomogeneous() { MetadataImage metadataImage2 = new MetadataImageBuilder() .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 2) .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 3) - .addRacks() .build(); Map members2 = new HashMap<>(); @@ -612,7 +605,6 @@ public void testAssignWithCurrentAssignmentHeterogeneous() { .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3) .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3) .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2) - .addRacks() .build(); Set memberATopicsSubscription1 = new LinkedHashSet<>(); @@ -685,7 +677,6 @@ public void testAssignWithCurrentAssignmentHeterogeneous() { .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3) .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2) .addTopic(TOPIC_4_UUID, TOPIC_4_NAME, 1) - .addRacks() .build(); Map members2 = new HashMap<>(); diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java index 7e6be4a5a3392..2c12ed467f17c 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java @@ -92,7 +92,6 @@ public Collection memberIds() { public void testTwoMembersNoTopicSubscription() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) - .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage @@ -130,7 +129,6 @@ public void testTwoMembersNoTopicSubscription() { public void testTwoMembersSubscribedToNonexistentTopics() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) - .addRacks() .build(); SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl( metadataImage @@ -166,7 +164,6 @@ public void testFirstAssignmentTwoMembersTwoTopics() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic3Uuid, topic3Name, 6) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -216,7 +213,6 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 2) .addTopic(topic3Uuid, topic3Name, 1) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -277,7 +273,6 @@ public void testReassignmentForTwoMembersThreeTopicsGivenUnbalancedPrevAssignmen .addTopic(topic1Uuid, topic1Name, 6) .addTopic(topic2Uuid, topic2Name, 4) .addTopic(topic3Uuid, topic3Name, 4) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -350,7 +345,6 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembers() { .addTopic(topic2Uuid, topic2Name, 5) .addTopic(topic3Uuid, topic3Name, 3) .addTopic(topic4Uuid, topic4Name, 3) - .addRacks() .build(); @@ -410,7 +404,6 @@ public void testReassignmentWhenOneMemberAddedAndPartitionsAddedTwoMembersTwoTop MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 6) .addTopic(topic2Uuid, topic2Name, 7) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -478,7 +471,6 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic2Uuid, topic2Name, 8) .addTopic(topic3Uuid, topic3Name, 3) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -535,7 +527,6 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) .addTopic(topic2Uuid, topic2Name, 5) - .addRacks() .build(); // Initial subscriptions were [T1, T2] @@ -598,7 +589,6 @@ public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions( .addTopic(topic1Uuid, topic1Name, 2) .addTopic(topic2Uuid, topic2Name, 2) .addTopic(topic3Uuid, topic3Name, 2) - .addRacks() .build(); Map members = new TreeMap<>(); @@ -652,7 +642,6 @@ public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions( public void testFirstAssignmentWithTwoMembersIncludingOneWithoutSubscriptions() { MetadataImage metadataImage = new MetadataImageBuilder() .addTopic(topic1Uuid, topic1Name, 3) - .addRacks() .build(); Map members = new TreeMap<>(); diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java index 47a8e43ff5117..b2ecfcbce78d4 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java @@ -26,7 +26,6 @@ import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.stream.IntStream; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; @@ -41,12 +40,11 @@ public class SubscribedTopicMetadataTest { @BeforeEach public void setUp() { MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder(); - IntStream.range(0, 5).forEach(i -> { + for (int i = 0; i < 5; i++) { Uuid topicId = Uuid.randomUuid(); String topicName = "topic" + i; metadataImageBuilder.addTopic(topicId, topicName, numPartitions); - }); - metadataImageBuilder.addRacks(); + } metadataImage = metadataImageBuilder.addRacks().build(); subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage); @@ -57,6 +55,11 @@ public void testMetadataImageCannotBeNull() { assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(null)); } + @Test + public void testTopicPartitionAllowedMapCannotBeNull() { + assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(metadataImage, null)); + } + @Test public void testNumberOfPartitions() { Uuid topicId = Uuid.randomUuid(); @@ -102,7 +105,7 @@ public void testEquals() { public void testAssignablePartitions() { String t1Name = "t1"; Uuid t1Id = Uuid.randomUuid(); - metadataImage = new MetadataImageBuilder().addTopic(t1Id, t1Name, numPartitions).addRacks().build(); + metadataImage = new MetadataImageBuilder().addTopic(t1Id, t1Name, numPartitions).build(); // Optional.empty() allow map (all partitions assignable) subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage, Optional.empty()); assertEquals(Set.of(0, 1, 2, 3, 4), subscribedTopicMetadata.assignablePartitions(t1Id)); From 49e215c5702379a2638df77daf21ab61cad029fa Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Mon, 19 May 2025 09:20:29 +0800 Subject: [PATCH 6/8] remove Set.copyOf from assignablePartitions Signed-off-by: PoAn Yang --- .../coordinator/group/modern/SubscribedTopicDescriberImpl.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java index 740769baf4198..384beb5d6843b 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java @@ -112,7 +112,7 @@ public Set assignablePartitions(Uuid topicId) { } if (topicPartitionAllowedMap.isEmpty()) { - return Set.copyOf(topic.partitions().keySet()); + return topic.partitions().keySet(); } return topicPartitionAllowedMap.get().getOrDefault(topicId, Set.of()); From ce3701aa7e588f16c552e1a81c08e40dff89ccdf Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Mon, 19 May 2025 15:07:23 +0800 Subject: [PATCH 7/8] address comment Signed-off-by: PoAn Yang --- .../coordinator/group/modern/SubscribedTopicDescriberImpl.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java index 384beb5d6843b..6a2aa02b99f63 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java @@ -23,6 +23,7 @@ import org.apache.kafka.image.TopicImage; import org.apache.kafka.metadata.PartitionRegistration; +import java.util.Collections; import java.util.HashSet; import java.util.Map; import java.util.Objects; @@ -112,7 +113,7 @@ public Set assignablePartitions(Uuid topicId) { } if (topicPartitionAllowedMap.isEmpty()) { - return topic.partitions().keySet(); + return Collections.unmodifiableSet(topic.partitions().keySet()); } return topicPartitionAllowedMap.get().getOrDefault(topicId, Set.of()); From 8076eed18bdb87c1548c12ad1c85b71ca8e22db5 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Mon, 19 May 2025 17:43:22 +0800 Subject: [PATCH 8/8] check whether BrokerRegistration is null Signed-off-by: PoAn Yang --- .../group/modern/SubscribedTopicDescriberImpl.java | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java index 6a2aa02b99f63..d8237724b748b 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java @@ -21,6 +21,7 @@ import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber; import org.apache.kafka.image.MetadataImage; import org.apache.kafka.image.TopicImage; +import org.apache.kafka.metadata.BrokerRegistration; import org.apache.kafka.metadata.PartitionRegistration; import java.util.Collections; @@ -88,9 +89,12 @@ public Set racksForPartition(Uuid topicId, int partition) { Set racks = new HashSet<>(); for (int replica : partitionRegistration.replicas) { // Only add the rack if it is available for the broker/replica. - metadataImage.cluster().broker(replica).rack().ifPresent(racks::add); + BrokerRegistration brokerRegistration = metadataImage.cluster().broker(replica); + if (brokerRegistration != null) { + brokerRegistration.rack().ifPresent(racks::add); + } } - return racks; + return Collections.unmodifiableSet(racks); } } return Set.of();