rreddy-22 commented on code in PR #15717: URL: https://github.com/apache/kafka/pull/15717#discussion_r1571435096
########## jmh-benchmarks/src/main/java/org/apache/kafka/jmh/group_coordinator/ServerSideAssignorBenchmark.java: ########## @@ -0,0 +1,185 @@ +package org.apache.kafka.jmh.group_coordinator; + +import org.apache.kafka.common.Uuid; +import org.apache.kafka.coordinator.group.assignor.AssignmentMemberSpec; +import org.apache.kafka.coordinator.group.assignor.AssignmentSpec; +import org.apache.kafka.coordinator.group.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.assignor.MemberAssignment; +import org.apache.kafka.coordinator.group.assignor.PartitionAssignor; +import org.apache.kafka.coordinator.group.assignor.RangeAssignor; +import org.apache.kafka.coordinator.group.assignor.SubscribedTopicDescriber; +import org.apache.kafka.coordinator.group.assignor.UniformAssignor; +import org.apache.kafka.coordinator.group.consumer.SubscribedTopicMetadata; +import org.apache.kafka.coordinator.group.consumer.TopicMetadata; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.annotations.Warmup; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.TreeMap; +import java.util.concurrent.TimeUnit; + +import static java.lang.Integer.max; + +@State(Scope.Benchmark) +@Fork(value = 1) +@Warmup(iterations = 5) +@Measurement(iterations = 5) +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +public class ServerSideAssignorBenchmark { + + @Param({"10", "50", "100"}) + private int partitionsPerTopicCount; + + @Param({"100"}) + private int topicCount; + + @Param({"500", "1000"}) + private int memberCount; + + @Param({"true", "false"}) + private boolean isRackAware; + + @Param({"true", "false"}) + private boolean isSubscriptionUniform; + + @Param({"true", "false"}) + private boolean isRangeAssignor; + + @Param({"true", "false"}) + private boolean isReassignment; + + private PartitionAssignor partitionAssignor; + + private final int numberOfRacks = 3; + + private AssignmentSpec assignmentSpec; + + private SubscribedTopicDescriber subscribedTopicDescriber; + + @Setup(Level.Trial) + public void setup() { + Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>(); + Map<Integer, Set<String>> partitionRacks = isRackAware ? + mkMapOfPartitionRacks(partitionsPerTopicCount) : + Collections.emptyMap(); + + for (int i = 1; i <= topicCount; i++) { + Uuid topicUuid = Uuid.randomUuid(); + String topicName = "topic" + i; + topicMetadata.put(topicUuid, new TopicMetadata( + topicUuid, topicName, partitionsPerTopicCount, partitionRacks)); + } + + addTopicSubscriptions(topicMetadata); + this.subscribedTopicDescriber = new SubscribedTopicMetadata(topicMetadata); + + if (isRangeAssignor) { + this.partitionAssignor = new RangeAssignor(); + } else { + this.partitionAssignor = new UniformAssignor(); + } + + if (isReassignment) { + GroupAssignment initialAssignment = partitionAssignor.assign(assignmentSpec, subscribedTopicDescriber); + Map<String, MemberAssignment> members; + + members = initialAssignment.members(); + + // Update the AssignmentSpec with the results from the initial assignment. + Map<String, AssignmentMemberSpec> updatedMembers = new HashMap<>(); + + members.forEach((memberId, memberAssignment) -> { + AssignmentMemberSpec memberSpec = assignmentSpec.members().get(memberId); + updatedMembers.put(memberId, new AssignmentMemberSpec( + memberSpec.instanceId(), + memberSpec.rackId(), + memberSpec.subscribedTopicIds(), + memberAssignment.targetPartitions() + )); + }); + + // Add new member to trigger a reassignment. + Optional<String> rackId = isRackAware ? Optional.of("rack" + (memberCount + 1) % numberOfRacks) : Optional.empty(); + + updatedMembers.put("newMember", new AssignmentMemberSpec( + Optional.empty(), + rackId, + topicMetadata.keySet(), + Collections.emptyMap() + )); + + this.assignmentSpec = new AssignmentSpec(updatedMembers); + } + } + + private Map<Integer, Set<String>> mkMapOfPartitionRacks(int numPartitions) { + Map<Integer, Set<String>> partitionRacks = new HashMap<>(numPartitions); + for (int i = 0; i < numPartitions; i++) { + partitionRacks.put(i, new HashSet<>(Arrays.asList("rack" + i % numberOfRacks, "rack" + (i + 1) % numberOfRacks))); + } + return partitionRacks; + } + + private void addTopicSubscriptions(Map<Uuid, TopicMetadata> topicMetadata) { + Map<String, AssignmentMemberSpec> members = new TreeMap<>(); + List<Uuid> allTopicIds = new ArrayList<>(topicMetadata.keySet()); + int topicCounter = 0; + + for (int i = 0; i < memberCount; i++) { + String memberName = "member" + i; + Optional<String> rackId = isRackAware ? Optional.of("rack" + i % numberOfRacks) : Optional.empty(); + List<Uuid> subscribedTopicIds; + + // When subscriptions are uniform, all members are assigned all topics. + if (isSubscriptionUniform) { + subscribedTopicIds = allTopicIds; + } else { + subscribedTopicIds = Arrays.asList( + allTopicIds.get(i % topicCount), + allTopicIds.get((i+1) % topicCount) + ); + topicCounter = max (topicCounter, ((i+1) % topicCount)); + + if (i == memberCount - 1 && topicCounter < topicCount - 1) { + subscribedTopicIds.addAll(allTopicIds.subList(topicCounter + 1, topicCount - 1)); + } Review Comment: yeah makes sense, I implemented it, but my only worry is that the max topics any member can get will be topicCount/5 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: jira-unsubscr...@kafka.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org