mirror of https://github.com/apache/kafka.git
KAFKA-16568: JMH Benchmarks for Server Side Rebalances (#15717)
This patch add three benchmarks for the client assignors, the server assignors and the target assignment builder. Reviewers: David Jacot <djacot@confluent.io>
This commit is contained in:
Ritika Reddy
GitHub
parent
0a6d5ff23c
commit
8013657f5d
|
|
@ -0,0 +1,287 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.jmh.assignor;
|
||||
|
||||
import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
|
||||
import org.apache.kafka.clients.consumer.CooperativeStickyAssignor;
|
||||
import org.apache.kafka.clients.consumer.RangeAssignor;
|
||||
import org.apache.kafka.common.Cluster;
|
||||
import org.apache.kafka.common.Node;
|
||||
import org.apache.kafka.common.PartitionInfo;
|
||||
import org.apache.kafka.common.TopicPartition;
|
||||
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.Collections;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
import static org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor.DEFAULT_GENERATION;
|
||||
|
||||
@State(Scope.Benchmark)
|
||||
@Fork(value = 1)
|
||||
@Warmup(iterations = 5)
|
||||
@Measurement(iterations = 5)
|
||||
@BenchmarkMode(Mode.AverageTime)
|
||||
@OutputTimeUnit(TimeUnit.MILLISECONDS)
|
||||
public class ClientSideAssignorBenchmark {
|
||||
|
||||
public enum AssignorType {
|
||||
RANGE(new RangeAssignor()),
|
||||
COOPERATIVE_STICKY(new CooperativeStickyAssignor());
|
||||
|
||||
private final ConsumerPartitionAssignor assignor;
|
||||
|
||||
AssignorType(ConsumerPartitionAssignor assignor) {
|
||||
this.assignor = assignor;
|
||||
}
|
||||
|
||||
public ConsumerPartitionAssignor assignor() {
|
||||
return assignor;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* The subscription pattern followed by the members of the group.
|
||||
*
|
||||
* A subscription model is considered homogenous if all the members of the group
|
||||
* are subscribed to the same set of topics, it is heterogeneous otherwise.
|
||||
*/
|
||||
public enum SubscriptionModel {
|
||||
HOMOGENEOUS, HETEROGENEOUS
|
||||
}
|
||||
|
||||
/**
|
||||
* The assignment type is decided based on whether all the members are assigned partitions
|
||||
* for the first time (full), or incrementally when a rebalance is triggered.
|
||||
*/
|
||||
public enum AssignmentType {
|
||||
FULL, INCREMENTAL
|
||||
}
|
||||
|
||||
@Param({"100", "500", "1000", "5000", "10000"})
|
||||
private int memberCount;
|
||||
|
||||
@Param({"5", "10", "50"})
|
||||
private int partitionsToMemberRatio;
|
||||
|
||||
@Param({"10", "100", "1000"})
|
||||
private int topicCount;
|
||||
|
||||
@Param({"true", "false"})
|
||||
private boolean isRackAware;
|
||||
|
||||
@Param({"HOMOGENEOUS", "HETEROGENEOUS"})
|
||||
private SubscriptionModel subscriptionModel;
|
||||
|
||||
@Param({"RANGE", "COOPERATIVE_STICKY"})
|
||||
private AssignorType assignorType;
|
||||
|
||||
@Param({"FULL", "INCREMENTAL"})
|
||||
private AssignmentType assignmentType;
|
||||
|
||||
private Map<String, ConsumerPartitionAssignor.Subscription> subscriptions = new HashMap<>();
|
||||
|
||||
private ConsumerPartitionAssignor.GroupSubscription groupSubscription;
|
||||
|
||||
private static final int NUMBER_OF_RACKS = 3;
|
||||
|
||||
private static final int MAX_BUCKET_COUNT = 5;
|
||||
|
||||
private ConsumerPartitionAssignor assignor;
|
||||
|
||||
private Cluster metadata;
|
||||
|
||||
private final List<String> allTopicNames = new ArrayList<>();
|
||||
|
||||
@Setup(Level.Trial)
|
||||
public void setup() {
|
||||
// Ensure there are enough racks and brokers for the replication factor = 3.
|
||||
if (NUMBER_OF_RACKS < 3) {
|
||||
throw new IllegalArgumentException("Number of broker racks must be at least equal to 3.");
|
||||
}
|
||||
|
||||
populateTopicMetadata();
|
||||
|
||||
addMemberSubscriptions();
|
||||
|
||||
assignor = assignorType.assignor();
|
||||
|
||||
if (assignmentType == AssignmentType.INCREMENTAL) {
|
||||
simulateIncrementalRebalance();
|
||||
}
|
||||
}
|
||||
|
||||
private void populateTopicMetadata() {
|
||||
List<PartitionInfo> partitions = new ArrayList<>();
|
||||
int partitionsPerTopicCount = (memberCount * partitionsToMemberRatio) / topicCount;
|
||||
|
||||
// Create nodes (brokers), one for each rack.
|
||||
List<Node> nodes = new ArrayList<>(NUMBER_OF_RACKS);
|
||||
for (int i = 0; i < NUMBER_OF_RACKS; i++) {
|
||||
nodes.add(new Node(i, "", i, "rack" + i));
|
||||
}
|
||||
|
||||
for (int i = 0; i < topicCount; i++) {
|
||||
String topicName = "topic" + i;
|
||||
allTopicNames.add(topicName);
|
||||
partitions.addAll(partitionInfos(topicName, partitionsPerTopicCount, nodes));
|
||||
}
|
||||
|
||||
metadata = new Cluster("test-cluster", nodes, partitions, Collections.emptySet(), Collections.emptySet());
|
||||
}
|
||||
|
||||
private void addMemberSubscriptions() {
|
||||
// In the rebalance case, we will add the last member as a trigger.
|
||||
// This is done to keep the final members count consistent with the input.
|
||||
int numberOfMembers = assignmentType.equals(AssignmentType.INCREMENTAL) ? memberCount - 1 : memberCount;
|
||||
|
||||
// When subscriptions are homogeneous, all members are assigned all topics.
|
||||
if (subscriptionModel == SubscriptionModel.HOMOGENEOUS) {
|
||||
for (int i = 0; i < numberOfMembers; i++) {
|
||||
String memberName = "member" + i;
|
||||
subscriptions.put(memberName, subscription(allTopicNames, i));
|
||||
}
|
||||
} else {
|
||||
// Adjust bucket count based on member count when member count < max bucket count.
|
||||
int bucketCount = Math.min(MAX_BUCKET_COUNT, numberOfMembers);
|
||||
|
||||
// Check minimum topics requirement
|
||||
if (topicCount < bucketCount) {
|
||||
throw new IllegalArgumentException("At least " + bucketCount + " topics are recommended for effective bucketing.");
|
||||
}
|
||||
|
||||
int bucketSizeTopics = (int) Math.ceil((double) topicCount / bucketCount);
|
||||
int bucketSizeMembers = (int) Math.ceil((double) numberOfMembers / bucketCount);
|
||||
|
||||
// Define buckets for each member and assign topics from the same bucket
|
||||
for (int bucket = 0; bucket < bucketCount; bucket++) {
|
||||
int memberStartIndex = bucket * bucketSizeMembers;
|
||||
int memberEndIndex = Math.min((bucket + 1) * bucketSizeMembers, numberOfMembers);
|
||||
|
||||
int topicStartIndex = bucket * bucketSizeTopics;
|
||||
int topicEndIndex = Math.min((bucket + 1) * bucketSizeTopics, topicCount);
|
||||
|
||||
List<String> bucketTopics = allTopicNames.subList(topicStartIndex, topicEndIndex);
|
||||
|
||||
// Assign topics to each member in the current bucket
|
||||
for (int i = memberStartIndex; i < memberEndIndex; i++) {
|
||||
String memberName = "member" + i;
|
||||
subscriptions.put(memberName, subscription(bucketTopics, i));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
groupSubscription = new ConsumerPartitionAssignor.GroupSubscription(subscriptions);
|
||||
}
|
||||
|
||||
private List<PartitionInfo> partitionInfos(String topic, int numberOfPartitions, List<Node> nodes) {
|
||||
// Create PartitionInfo for each partition.
|
||||
// Replication factor is hardcoded to 2.
|
||||
List<PartitionInfo> partitionInfos = new ArrayList<>(numberOfPartitions);
|
||||
for (int i = 0; i < numberOfPartitions; i++) {
|
||||
Node[] replicas = new Node[3];
|
||||
for (int j = 0; j < 3; j++) {
|
||||
// Assign nodes based on partition number to mimic mkMapOfPartitionRacks logic.
|
||||
int nodeIndex = (i + j) % NUMBER_OF_RACKS;
|
||||
replicas[j] = nodes.get(nodeIndex);
|
||||
}
|
||||
partitionInfos.add(new PartitionInfo(topic, i, replicas[0], replicas, replicas));
|
||||
}
|
||||
|
||||
return partitionInfos;
|
||||
}
|
||||
|
||||
private ConsumerPartitionAssignor.Subscription subscription(List<String> topics, int consumerIndex) {
|
||||
Optional<String> rackId = rackId(consumerIndex);
|
||||
return new ConsumerPartitionAssignor.Subscription(
|
||||
topics,
|
||||
null,
|
||||
Collections.emptyList(),
|
||||
DEFAULT_GENERATION,
|
||||
rackId
|
||||
);
|
||||
}
|
||||
|
||||
private Optional<String> rackId(int consumerIndex) {
|
||||
return isRackAware ? Optional.of("rack" + consumerIndex % NUMBER_OF_RACKS) : Optional.empty();
|
||||
}
|
||||
|
||||
private ConsumerPartitionAssignor.Subscription subscriptionWithOwnedPartitions(
|
||||
List<TopicPartition> ownedPartitions,
|
||||
ConsumerPartitionAssignor.Subscription prevSubscription
|
||||
) {
|
||||
return new ConsumerPartitionAssignor.Subscription(
|
||||
prevSubscription.topics(),
|
||||
null,
|
||||
ownedPartitions,
|
||||
DEFAULT_GENERATION,
|
||||
prevSubscription.rackId()
|
||||
);
|
||||
}
|
||||
|
||||
private void simulateIncrementalRebalance() {
|
||||
ConsumerPartitionAssignor.GroupAssignment initialAssignment = assignor.assign(metadata, groupSubscription);
|
||||
Map<String, ConsumerPartitionAssignor.Subscription> newSubscriptions = new HashMap<>();
|
||||
subscriptions.forEach((member, subscription) ->
|
||||
newSubscriptions.put(
|
||||
member,
|
||||
subscriptionWithOwnedPartitions(
|
||||
initialAssignment.groupAssignment().get(member).partitions(),
|
||||
subscription
|
||||
)
|
||||
)
|
||||
);
|
||||
|
||||
List<String> subscribedTopicsForNewMember;
|
||||
if (subscriptionModel == SubscriptionModel.HETEROGENEOUS) {
|
||||
subscribedTopicsForNewMember = subscriptions.get("member" + (memberCount - 2)).topics();
|
||||
} else {
|
||||
subscribedTopicsForNewMember = allTopicNames;
|
||||
}
|
||||
|
||||
// Add new member to trigger a reassignment.
|
||||
newSubscriptions.put("newMember", subscription(
|
||||
subscribedTopicsForNewMember,
|
||||
memberCount - 1
|
||||
));
|
||||
|
||||
subscriptions = newSubscriptions;
|
||||
}
|
||||
|
||||
@Benchmark
|
||||
@Threads(1)
|
||||
@OutputTimeUnit(TimeUnit.MILLISECONDS)
|
||||
public void doAssignment() {
|
||||
assignor.assign(metadata, groupSubscription);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,281 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.jmh.assignor;
|
||||
|
||||
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.Collection;
|
||||
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.concurrent.TimeUnit;
|
||||
|
||||
@State(Scope.Benchmark)
|
||||
@Fork(value = 1)
|
||||
@Warmup(iterations = 5)
|
||||
@Measurement(iterations = 5)
|
||||
@BenchmarkMode(Mode.AverageTime)
|
||||
@OutputTimeUnit(TimeUnit.MILLISECONDS)
|
||||
public class ServerSideAssignorBenchmark {
|
||||
|
||||
public enum AssignorType {
|
||||
RANGE(new RangeAssignor()),
|
||||
UNIFORM(new UniformAssignor());
|
||||
|
||||
private final PartitionAssignor assignor;
|
||||
|
||||
AssignorType(PartitionAssignor assignor) {
|
||||
this.assignor = assignor;
|
||||
}
|
||||
|
||||
public PartitionAssignor assignor() {
|
||||
return assignor;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* The subscription pattern followed by the members of the group.
|
||||
*
|
||||
* A subscription model is considered homogenous if all the members of the group
|
||||
* are subscribed to the same set of topics, it is heterogeneous otherwise.
|
||||
*/
|
||||
public enum SubscriptionModel {
|
||||
HOMOGENEOUS, HETEROGENEOUS
|
||||
}
|
||||
|
||||
/**
|
||||
* The assignment type is decided based on whether all the members are assigned partitions
|
||||
* for the first time (full), or incrementally when a rebalance is triggered.
|
||||
*/
|
||||
public enum AssignmentType {
|
||||
FULL, INCREMENTAL
|
||||
}
|
||||
|
||||
@Param({"100", "500", "1000", "5000", "10000"})
|
||||
private int memberCount;
|
||||
|
||||
@Param({"5", "10", "50"})
|
||||
private int partitionsToMemberRatio;
|
||||
|
||||
@Param({"10", "100", "1000"})
|
||||
private int topicCount;
|
||||
|
||||
@Param({"true", "false"})
|
||||
private boolean isRackAware;
|
||||
|
||||
@Param({"HOMOGENEOUS", "HETEROGENEOUS"})
|
||||
private SubscriptionModel subscriptionModel;
|
||||
|
||||
@Param({"RANGE", "UNIFORM"})
|
||||
private AssignorType assignorType;
|
||||
|
||||
@Param({"FULL", "INCREMENTAL"})
|
||||
private AssignmentType assignmentType;
|
||||
|
||||
private PartitionAssignor partitionAssignor;
|
||||
|
||||
private static final int NUMBER_OF_RACKS = 3;
|
||||
|
||||
private static final int MAX_BUCKET_COUNT = 5;
|
||||
|
||||
private AssignmentSpec assignmentSpec;
|
||||
|
||||
private SubscribedTopicDescriber subscribedTopicDescriber;
|
||||
|
||||
private final List<Uuid> allTopicIds = new ArrayList<>();
|
||||
|
||||
@Setup(Level.Trial)
|
||||
public void setup() {
|
||||
Map<Uuid, TopicMetadata> topicMetadata = createTopicMetadata();
|
||||
subscribedTopicDescriber = new SubscribedTopicMetadata(topicMetadata);
|
||||
|
||||
createAssignmentSpec();
|
||||
|
||||
partitionAssignor = assignorType.assignor();
|
||||
|
||||
if (assignmentType == AssignmentType.INCREMENTAL) {
|
||||
simulateIncrementalRebalance();
|
||||
}
|
||||
}
|
||||
|
||||
private Map<Uuid, TopicMetadata> createTopicMetadata() {
|
||||
Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
|
||||
int partitionsPerTopicCount = (memberCount * partitionsToMemberRatio) / topicCount;
|
||||
|
||||
Map<Integer, Set<String>> partitionRacks = isRackAware ?
|
||||
mkMapOfPartitionRacks(partitionsPerTopicCount) :
|
||||
Collections.emptyMap();
|
||||
|
||||
for (int i = 0; i < topicCount; i++) {
|
||||
Uuid topicUuid = Uuid.randomUuid();
|
||||
String topicName = "topic" + i;
|
||||
allTopicIds.add(topicUuid);
|
||||
topicMetadata.put(topicUuid, new TopicMetadata(
|
||||
topicUuid,
|
||||
topicName,
|
||||
partitionsPerTopicCount,
|
||||
partitionRacks
|
||||
));
|
||||
}
|
||||
|
||||
return topicMetadata;
|
||||
}
|
||||
|
||||
private void createAssignmentSpec() {
|
||||
Map<String, AssignmentMemberSpec> members = new HashMap<>();
|
||||
|
||||
// In the rebalance case, we will add the last member as a trigger.
|
||||
// This is done to keep the total members count consistent with the input.
|
||||
int numberOfMembers = assignmentType.equals(AssignmentType.INCREMENTAL) ? memberCount - 1 : memberCount;
|
||||
|
||||
if (subscriptionModel.equals(SubscriptionModel.HOMOGENEOUS)) {
|
||||
for (int i = 0; i < numberOfMembers; i++) {
|
||||
addMemberSpec(members, i, new HashSet<>(allTopicIds));
|
||||
}
|
||||
} else {
|
||||
// Adjust bucket count based on member count when member count < max bucket count.
|
||||
int bucketCount = Math.min(MAX_BUCKET_COUNT, numberOfMembers);
|
||||
|
||||
// Check minimum topics requirement
|
||||
if (topicCount < bucketCount) {
|
||||
throw new IllegalArgumentException("At least " + bucketCount + " topics are recommended for effective bucketing.");
|
||||
}
|
||||
|
||||
int bucketSizeTopics = (int) Math.ceil((double) topicCount / bucketCount);
|
||||
int bucketSizeMembers = (int) Math.ceil((double) numberOfMembers / bucketCount);
|
||||
|
||||
// Define buckets for each member and assign topics from the same bucket
|
||||
for (int bucket = 0; bucket < bucketCount; bucket++) {
|
||||
int memberStartIndex = bucket * bucketSizeMembers;
|
||||
int memberEndIndex = Math.min((bucket + 1) * bucketSizeMembers, numberOfMembers);
|
||||
|
||||
int topicStartIndex = bucket * bucketSizeTopics;
|
||||
int topicEndIndex = Math.min((bucket + 1) * bucketSizeTopics, topicCount);
|
||||
|
||||
Set<Uuid> bucketTopics = new HashSet<>(allTopicIds.subList(topicStartIndex, topicEndIndex));
|
||||
|
||||
// Assign topics to each member in the current bucket
|
||||
for (int i = memberStartIndex; i < memberEndIndex; i++) {
|
||||
addMemberSpec(members, i, bucketTopics);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
this.assignmentSpec = new AssignmentSpec(members);
|
||||
}
|
||||
|
||||
private Optional<String> rackId(int memberIndex) {
|
||||
return isRackAware ? Optional.of("rack" + memberIndex % NUMBER_OF_RACKS) : Optional.empty();
|
||||
}
|
||||
|
||||
private void addMemberSpec(
|
||||
Map<String, AssignmentMemberSpec> members,
|
||||
int memberIndex,
|
||||
Set<Uuid> subscribedTopicIds
|
||||
) {
|
||||
String memberId = "member" + memberIndex;
|
||||
Optional<String> rackId = rackId(memberIndex);
|
||||
|
||||
members.put(memberId, new AssignmentMemberSpec(
|
||||
Optional.empty(),
|
||||
rackId,
|
||||
subscribedTopicIds,
|
||||
Collections.emptyMap()
|
||||
));
|
||||
}
|
||||
|
||||
private static 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 % NUMBER_OF_RACKS,
|
||||
"rack" + (i + 1) % NUMBER_OF_RACKS,
|
||||
"rack" + (i + 2) % NUMBER_OF_RACKS
|
||||
)));
|
||||
}
|
||||
return partitionRacks;
|
||||
}
|
||||
|
||||
private void simulateIncrementalRebalance() {
|
||||
GroupAssignment initialAssignment = partitionAssignor.assign(assignmentSpec, subscribedTopicDescriber);
|
||||
Map<String, MemberAssignment> members = initialAssignment.members();
|
||||
|
||||
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()
|
||||
));
|
||||
});
|
||||
|
||||
Collection<Uuid> subscribedTopicIdsForNewMember;
|
||||
if (subscriptionModel == SubscriptionModel.HETEROGENEOUS) {
|
||||
subscribedTopicIdsForNewMember = updatedMembers.get("member" + (memberCount - 2)).subscribedTopicIds();
|
||||
} else {
|
||||
subscribedTopicIdsForNewMember = allTopicIds;
|
||||
}
|
||||
|
||||
Optional<String> rackId = rackId(memberCount - 1);
|
||||
updatedMembers.put("newMember", new AssignmentMemberSpec(
|
||||
Optional.empty(),
|
||||
rackId,
|
||||
subscribedTopicIdsForNewMember,
|
||||
Collections.emptyMap()
|
||||
));
|
||||
|
||||
assignmentSpec = new AssignmentSpec(updatedMembers);
|
||||
}
|
||||
|
||||
@Benchmark
|
||||
@Threads(1)
|
||||
@OutputTimeUnit(TimeUnit.MILLISECONDS)
|
||||
public void doAssignment() {
|
||||
partitionAssignor.assign(assignmentSpec, subscribedTopicDescriber);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,195 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.jmh.assignor;
|
||||
|
||||
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.UniformAssignor;
|
||||
import org.apache.kafka.coordinator.group.consumer.Assignment;
|
||||
import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
|
||||
import org.apache.kafka.coordinator.group.consumer.SubscribedTopicMetadata;
|
||||
import org.apache.kafka.coordinator.group.consumer.TargetAssignmentBuilder;
|
||||
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.Collections;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
import java.util.Set;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
@State(Scope.Benchmark)
|
||||
@Fork(value = 1)
|
||||
@Warmup(iterations = 5)
|
||||
@Measurement(iterations = 5)
|
||||
@BenchmarkMode(Mode.AverageTime)
|
||||
@OutputTimeUnit(TimeUnit.MILLISECONDS)
|
||||
public class TargetAssignmentBuilderBenchmark {
|
||||
|
||||
@Param({"100", "500", "1000", "5000", "10000"})
|
||||
private int memberCount;
|
||||
|
||||
@Param({"5", "10", "50"})
|
||||
private int partitionsToMemberRatio;
|
||||
|
||||
@Param({"10", "100", "1000"})
|
||||
private int topicCount;
|
||||
|
||||
private static final String GROUP_ID = "benchmark-group";
|
||||
|
||||
private static final int GROUP_EPOCH = 0;
|
||||
|
||||
private PartitionAssignor partitionAssignor;
|
||||
|
||||
private Map<String, TopicMetadata> subscriptionMetadata = Collections.emptyMap();
|
||||
|
||||
private TargetAssignmentBuilder targetAssignmentBuilder;
|
||||
|
||||
private AssignmentSpec assignmentSpec;
|
||||
|
||||
private final List<String> allTopicNames = new ArrayList<>();
|
||||
|
||||
private final List<Uuid> allTopicIds = new ArrayList<>();
|
||||
|
||||
@Setup(Level.Trial)
|
||||
public void setup() {
|
||||
// For this benchmark we will use the Uniform Assignor
|
||||
// and a group that has a homogeneous subscription model.
|
||||
partitionAssignor = new UniformAssignor();
|
||||
|
||||
subscriptionMetadata = generateMockSubscriptionMetadata();
|
||||
Map<String, ConsumerGroupMember> members = generateMockMembers();
|
||||
Map<String, Assignment> existingTargetAssignment = generateMockInitialTargetAssignment();
|
||||
|
||||
ConsumerGroupMember newMember = new ConsumerGroupMember.Builder("newMember")
|
||||
.setSubscribedTopicNames(allTopicNames)
|
||||
.build();
|
||||
|
||||
targetAssignmentBuilder = new TargetAssignmentBuilder(GROUP_ID, GROUP_EPOCH, partitionAssignor)
|
||||
.withMembers(members)
|
||||
.withSubscriptionMetadata(subscriptionMetadata)
|
||||
.withTargetAssignment(existingTargetAssignment)
|
||||
.addOrUpdateMember(newMember.memberId(), newMember);
|
||||
}
|
||||
|
||||
private Map<String, ConsumerGroupMember> generateMockMembers() {
|
||||
Map<String, ConsumerGroupMember> members = new HashMap<>();
|
||||
|
||||
for (int i = 0; i < memberCount - 1; i++) {
|
||||
ConsumerGroupMember member = new ConsumerGroupMember.Builder("member" + i)
|
||||
.setSubscribedTopicNames(allTopicNames)
|
||||
.build();
|
||||
members.put("member" + i, member);
|
||||
}
|
||||
return members;
|
||||
}
|
||||
|
||||
private Map<String, TopicMetadata> generateMockSubscriptionMetadata() {
|
||||
Map<String, TopicMetadata> subscriptionMetadata = new HashMap<>();
|
||||
int partitionsPerTopicCount = (memberCount * partitionsToMemberRatio) / topicCount;
|
||||
|
||||
for (int i = 0; i < topicCount; i++) {
|
||||
String topicName = "topic-" + i;
|
||||
Uuid topicId = Uuid.randomUuid();
|
||||
allTopicNames.add(topicName);
|
||||
allTopicIds.add(topicId);
|
||||
|
||||
TopicMetadata metadata = new TopicMetadata(
|
||||
topicId,
|
||||
topicName,
|
||||
partitionsPerTopicCount,
|
||||
Collections.emptyMap()
|
||||
);
|
||||
subscriptionMetadata.put(topicName, metadata);
|
||||
}
|
||||
|
||||
return subscriptionMetadata;
|
||||
}
|
||||
|
||||
private Map<String, Assignment> generateMockInitialTargetAssignment() {
|
||||
Map<Uuid, TopicMetadata> topicMetadataMap = new HashMap<>(topicCount);
|
||||
subscriptionMetadata.forEach((topicName, topicMetadata) ->
|
||||
topicMetadataMap.put(
|
||||
topicMetadata.id(),
|
||||
topicMetadata
|
||||
)
|
||||
);
|
||||
|
||||
createAssignmentSpec();
|
||||
|
||||
GroupAssignment groupAssignment = partitionAssignor.assign(
|
||||
assignmentSpec,
|
||||
new SubscribedTopicMetadata(topicMetadataMap)
|
||||
);
|
||||
|
||||
Map<String, Assignment> initialTargetAssignment = new HashMap<>(memberCount);
|
||||
|
||||
for (Map.Entry<String, MemberAssignment> entry : groupAssignment.members().entrySet()) {
|
||||
String memberId = entry.getKey();
|
||||
Map<Uuid, Set<Integer>> topicPartitions = entry.getValue().targetPartitions();
|
||||
|
||||
Assignment assignment = new Assignment(topicPartitions);
|
||||
|
||||
initialTargetAssignment.put(memberId, assignment);
|
||||
}
|
||||
|
||||
return initialTargetAssignment;
|
||||
}
|
||||
|
||||
private void createAssignmentSpec() {
|
||||
Map<String, AssignmentMemberSpec> members = new HashMap<>();
|
||||
|
||||
for (int i = 0; i < memberCount - 1; i++) {
|
||||
String memberId = "member" + i;
|
||||
|
||||
members.put(memberId, new AssignmentMemberSpec(
|
||||
Optional.empty(),
|
||||
Optional.empty(),
|
||||
allTopicIds,
|
||||
Collections.emptyMap()
|
||||
));
|
||||
}
|
||||
assignmentSpec = new AssignmentSpec(members);
|
||||
}
|
||||
|
||||
@Benchmark
|
||||
@Threads(1)
|
||||
@OutputTimeUnit(TimeUnit.MILLISECONDS)
|
||||
public void build() {
|
||||
targetAssignmentBuilder.build();
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue