(kafka) branch trunk updated: KAFKA-17496: Add heterogeneous case to TargetAssignmentBuilderBenchmark (#17277)

[dajac]

This is an automated email from the ASF dual-hosted git repository.

dajac pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git


The following commit(s) were added to refs/heads/trunk by this push:
     new 236f3d422f6 KAFKA-17496: Add heterogeneous case to 
TargetAssignmentBuilderBenchmark (#17277)
236f3d422f6 is described below

commit 236f3d422f659710f02120ac6cf3de229ee9a53f
Author: Sean Quah <sq...@confluent.io>
AuthorDate: Thu Sep 26 07:59:38 2024 +0100

    KAFKA-17496: Add heterogeneous case to TargetAssignmentBuilderBenchmark 
(#17277)
    
    Bring the homogeneous case from ServerSideAssignorBenchmark to 
TargetAssignmentBuilderBenchmark.
    
    Reviewers: David Jacot <dja...@confluent.io>
---
 .../jmh/assignor/ServerSideAssignorBenchmark.java  |  1 +
 .../assignor/TargetAssignmentBuilderBenchmark.java | 33 ++++++++++++++++------
 2 files changed, 26 insertions(+), 8 deletions(-)

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 a22673c3f59..1ad30e320b5 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
@@ -118,6 +118,7 @@ public class ServerSideAssignorBenchmark {
 
     private static final int NUMBER_OF_RACKS = 3;
 
+    /** The number of homogeneous subgroups to create for the heterogeneous 
subscription case. */
     private static final int MAX_BUCKET_COUNT = 5;
 
     private GroupSpec groupSpec;
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 c9ddb54363b..a06e9b0e527 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
@@ -22,6 +22,7 @@ import 
org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
 import org.apache.kafka.coordinator.group.api.assignor.MemberAssignment;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor;
 import 
org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType;
 import org.apache.kafka.coordinator.group.assignor.UniformAssignor;
 import org.apache.kafka.coordinator.group.modern.Assignment;
 import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
@@ -71,6 +72,9 @@ public class TargetAssignmentBuilderBenchmark {
     @Param({"10", "100", "1000"})
     private int topicCount;
 
+    @Param({"HOMOGENEOUS", "HETEROGENEOUS"})
+    private SubscriptionType subscriptionType;
+
     private static final String GROUP_ID = "benchmark-group";
 
     private static final int GROUP_EPOCH = 0;
@@ -79,6 +83,9 @@ public class TargetAssignmentBuilderBenchmark {
 
     private TargetAssignmentBuilder<ConsumerGroupMember> 
targetAssignmentBuilder;
 
+    /** The number of homogeneous subgroups to create for the heterogeneous 
subscription case. */
+    private static final int MAX_BUCKET_COUNT = 5;
+
     private GroupSpec groupSpec;
 
     private Map<Uuid, Map<Integer, String>> invertedTargetAssignment;
@@ -109,7 +116,7 @@ public class TargetAssignmentBuilderBenchmark {
         targetAssignmentBuilder = new 
TargetAssignmentBuilder<ConsumerGroupMember>(GROUP_ID, GROUP_EPOCH, 
partitionAssignor)
             .withMembers(members)
             .withSubscriptionMetadata(subscriptionMetadata)
-            .withSubscriptionType(HOMOGENEOUS)
+            .withSubscriptionType(subscriptionType)
             .withTargetAssignment(existingTargetAssignment)
             .withInvertedTargetAssignment(invertedTargetAssignment)
             .withTopicsImage(topicsImage)
@@ -136,7 +143,7 @@ public class TargetAssignmentBuilderBenchmark {
     ) {
         this.groupSpec = AssignorBenchmarkUtils.createGroupSpec(
             members,
-            HOMOGENEOUS,
+            subscriptionType,
             topicsImage
         );
 
@@ -158,12 +165,22 @@ public class TargetAssignmentBuilderBenchmark {
     }
 
     private Map<String, ConsumerGroupMember> createMembers() {
-        return AssignorBenchmarkUtils.createHomogeneousMembers(
-            memberCount - 1,
-            this::memberId,
-            this::rackId,
-            allTopicNames
-        );
+        if (subscriptionType == HOMOGENEOUS) {
+            return AssignorBenchmarkUtils.createHomogeneousMembers(
+                memberCount - 1,
+                this::memberId,
+                this::rackId,
+                allTopicNames
+            );
+        } else {
+            return AssignorBenchmarkUtils.createHeterogeneousBucketedMembers(
+                memberCount - 1,
+                MAX_BUCKET_COUNT,
+                this::memberId,
+                this::rackId,
+                allTopicNames
+            );
+        }
     }
 
     private String memberId(int memberIndex) {