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

chia7712 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 a3123b01eca KAFKA-967 Use key range in ProducerPerformance (#22326)
a3123b01eca is described below

commit a3123b01ecac52306d603c54099fcaa498b5bea8
Author: Ken Huang <[email protected]>
AuthorDate: Mon Jul 27 10:57:04 2026 +0800

    KAFKA-967 Use key range in ProducerPerformance (#22326)
    
    Currently, kafka-producer-perf-test.sh only produces records with null 
keys, which makes it impossible to benchmark real-world keyed workloads (e.g., 
compacted topics, semantic partitioning, and stream joins). This PR implements 
KIP-1299 to introduce key generation support for the producer performance tool.
    
    Reviewers: Chia-Ping Tsai <[email protected]>
---
 checkstyle/suppressions.xml                        |   2 +
 docs/getting-started/upgrade.md                    |   1 +
 .../apache/kafka/tools/ProducerPerformance.java    |  75 +++++++++-
 .../kafka/tools/ProducerPerformanceTest.java       | 156 +++++++++++++++++++++
 4 files changed, 232 insertions(+), 2 deletions(-)

diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index 3823ff81ffc..4f2c312f248 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -244,6 +244,8 @@
               
files="(ProduceBenchSpec|ConsumeBenchSpec|SustainedConnectionSpec).java"/>
     <suppress id="dontUseSystemExit"
               
files="(VerifiableConsumer|VerifiableProducer|VerifiableShareConsumer).java"/>
+    <suppress checks="MethodLength"
+              files="ProducerPerformance\.java"/>
 
     <!-- Shell -->
     <suppress checks="CyclomaticComplexity"
diff --git a/docs/getting-started/upgrade.md b/docs/getting-started/upgrade.md
index 657b165b6da..4e8a7700190 100644
--- a/docs/getting-started/upgrade.md
+++ b/docs/getting-started/upgrade.md
@@ -47,6 +47,7 @@ type: docs
   * When clients connect to the cluster, they now include cluster and node 
information to enable detection and handling of misrouted connections. For 
further details, please refer to 
[KIP-1242](https://cwiki.apache.org/confluence/x/W4LMFw).
   * The `kafka-cluster.sh` tool now provides an `api-versions` command to 
display the API versions supported by the brokers or controllers, and it 
accepts both `--bootstrap-server` and `--bootstrap-controller`. As a result, 
`kafka-broker-api-versions.sh` is deprecated and will be removed in the next 
major release; use `kafka-cluster.sh api-versions` instead. For further 
details, please refer to 
[KIP-1220](https://cwiki.apache.org/confluence/x/-QkbFw).
   * Brokers can now record a human-readable description of each streams 
group's processing topology via a pluggable backend, retrievable through 
`Admin#describeStreamsGroups` and `kafka-streams-groups.sh --describe 
--topology`. The feature is disabled unless the new broker configuration 
`group.streams.topology.description.plugin.class` is set to a 
`StreamsGroupTopologyDescriptionPlugin` implementation; on the client side, the 
new Kafka Streams configuration `topology.description.push.ena [...]
+  * The `kafka-producer-perf-test.sh` tool now supports `--record-key-range`, 
`--key-distribution`, and `--random-seed` options to control the distribution 
of record keys. Use `--key-distribution range` for sequential key assignment 
(round-robin over the key range) or `--key-distribution random` for random key 
selection. The `--random-seed` option allows reproducible benchmark runs when 
using random key distribution. For further details, please refer to 
[KIP-1299](https://cwiki.apache.or [...]
 
 ## Upgrading to 4.3.0
 
diff --git 
a/tools/src/main/java/org/apache/kafka/tools/ProducerPerformance.java 
b/tools/src/main/java/org/apache/kafka/tools/ProducerPerformance.java
index 2b36e960d7f..078c3c4ddec 100644
--- a/tools/src/main/java/org/apache/kafka/tools/ProducerPerformance.java
+++ b/tools/src/main/java/org/apache/kafka/tools/ProducerPerformance.java
@@ -40,6 +40,7 @@ import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Locale;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.Scanner;
@@ -53,6 +54,14 @@ public class ProducerPerformance {
     public static final String DEFAULT_TRANSACTION_ID_PREFIX = 
"performance-producer-";
     public static final long DEFAULT_TRANSACTION_DURATION_MS = 3000L;
 
+    public enum KeyDistribution {
+        NONE, RANGE, RANDOM;
+
+        public static KeyDistribution fromString(String value) {
+            return KeyDistribution.valueOf(value.toUpperCase(Locale.ROOT));
+        }
+    }
+
     public static void main(String[] args) throws Exception {
         ProducerPerformance perf = new ProducerPerformance();
         perf.start(args);
@@ -74,7 +83,7 @@ public class ProducerPerformance {
                 payload = new byte[config.recordSize];
             }
             // not thread-safe, do not share with other threads
-            SplittableRandom random = new SplittableRandom(0);
+            SplittableRandom random = new SplittableRandom(config.randomSeed);
             ProducerRecord<byte[], byte[]> record;
 
             if (config.warmupRecords > 0) {
@@ -97,7 +106,8 @@ public class ProducerPerformance {
                     transactionStartTime = System.currentTimeMillis();
                 }
 
-                record = new ProducerRecord<>(config.topicName, payload);
+                byte[] key = generateKey(config.keyDistribution, 
config.recordKeyRange, i, random);
+                record = new ProducerRecord<>(config.topicName, null, key, 
payload);
 
                 long sendStartMs = System.currentTimeMillis();
                 if ((isSteadyState = config.warmupRecords > 0) && i == 
config.warmupRecords) {
@@ -167,6 +177,19 @@ public class ProducerPerformance {
     Stats stats;
     Stats steadyStateStats;
 
+    static byte[] generateKey(
+        KeyDistribution keyDistribution,
+        Integer recordKeyRange,
+        long recordIndex,
+        SplittableRandom random
+    ) {
+        return switch (keyDistribution) {
+            case RANGE -> Integer.toString((int) (recordIndex % 
recordKeyRange)).getBytes(StandardCharsets.UTF_8);
+            case RANDOM -> 
Integer.toString(random.nextInt(recordKeyRange)).getBytes(StandardCharsets.UTF_8);
+            default -> null;
+        };
+    }
+
     static byte[] generateRandomPayload(Integer recordSize, List<byte[]> 
payloadByteList, byte[] payload,
             SplittableRandom random, boolean payloadMonotonic, long 
recordValue) {
         if (!payloadByteList.isEmpty()) {
@@ -395,6 +418,37 @@ public class ProducerPerformance {
                 .setDefault(5_000L)
                 .help("Interval in milliseconds at which to print progress 
info.");
 
+        parser.addArgument("--record-key-range")
+                .action(store())
+                .required(false)
+                .type(Integer.class)
+                .metavar("KEY-RANGE")
+                .dest("recordKeyRange")
+                .help("The range of keys to use when --key-distribution is 
'range' or 'random'. " +
+                        "Keys will be integers in [0, KEY-RANGE). Required for 
range and random distributions.");
+
+        parser.addArgument("--key-distribution")
+                .action(store())
+                .required(false)
+                .type(String.class)
+                .metavar("KEY-DISTRIBUTION")
+                .dest("keyDistribution")
+                .choices("none", "range", "random")
+                .setDefault("none")
+                .help("The key distribution to use: 'none' for null keys, 
'range' for round-robin keys in " +
+                        "[0, KEY-RANGE), or 'random' for random keys in [0, 
KEY-RANGE). " +
+                        "Requires --record-key-range when set to 'range' or 
'random'.");
+
+        parser.addArgument("--random-seed")
+                .action(store())
+                .required(false)
+                .type(Long.class)
+                .metavar("RANDOM-SEED")
+                .dest("randomSeed")
+                .setDefault(0L)
+                .help("Seed for the pseudo-random number generator used by 
--key-distribution random and " +
+                        "random payload generation. The default value of 0 
ensures deterministic, reproducible " +
+                        "benchmark runs. Set to a different value when 
non-repeating sequences are required.");
         return parser;
     }
 
@@ -579,6 +633,9 @@ public class ProducerPerformance {
         final boolean transactionsEnabled;
         final List<byte[]> payloadByteList;
         final long reportingInterval;
+        final Integer recordKeyRange;
+        final KeyDistribution keyDistribution;
+        final long randomSeed;
 
         public ConfigPostProcessor(ArgumentParser parser, String[] args) 
throws IOException, ArgumentParserException {
             Namespace namespace = parser.parseArgs(args);
@@ -623,6 +680,20 @@ public class ProducerPerformance {
             if (reportingInterval <= 0) {
                 throw new ArgumentParserException("--reporting-interval should 
be greater than zero.", parser);
             }
+            this.recordKeyRange = namespace.getInt("recordKeyRange");
+            if (recordKeyRange != null && recordKeyRange <= 0) {
+                throw new ArgumentParserException("--record-key-range should 
be greater than zero.", parser);
+            }
+            this.keyDistribution = 
KeyDistribution.fromString(namespace.getString("keyDistribution"));
+            if (this.keyDistribution != KeyDistribution.NONE && recordKeyRange 
== null) {
+                throw new ArgumentParserException(
+                        "--record-key-range is required when 
--key-distribution is 'range' or 'random'.", parser);
+            }
+            if (this.keyDistribution == KeyDistribution.NONE && recordKeyRange 
!= null) {
+                throw new ArgumentParserException(
+                        "--key-distribution must be 'range' or 'random' when 
--record-key-range is specified.", parser);
+            }
+            this.randomSeed = namespace.getLong("randomSeed");
 
             // since default value gets printed with the help text, we are 
escaping \n there and replacing it with correct value here.
             String payloadDelimiter = 
namespace.getString("payloadDelimiter").equals("\\n")
diff --git 
a/tools/src/test/java/org/apache/kafka/tools/ProducerPerformanceTest.java 
b/tools/src/test/java/org/apache/kafka/tools/ProducerPerformanceTest.java
index 902de6ec740..0419c8cb1f5 100644
--- a/tools/src/test/java/org/apache/kafka/tools/ProducerPerformanceTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/ProducerPerformanceTest.java
@@ -19,6 +19,7 @@ package org.apache.kafka.tools;
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.errors.AuthorizationException;
 import org.apache.kafka.common.utils.Utils;
 
@@ -48,6 +49,7 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -792,4 +794,158 @@ public class ProducerPerformanceTest {
         assertEquals("10", 
configs.producerProps.get(ProducerConfig.LINGER_MS_CONFIG));
         assertEquals("32768", 
configs.producerProps.get(ProducerConfig.BATCH_SIZE_CONFIG));
     }
+
+    @Test
+    public void testKeyDistributionNoneByDefault() throws IOException, 
ArgumentParserException {
+        ArgumentParser parser = ProducerPerformance.argParser();
+        String[] args = new String[]{
+            "--topic", "Hello-Kafka",
+            "--num-records", "5",
+            "--throughput", "100",
+            "--record-size", "100",
+            "--bootstrap-server", "localhost:9000"};
+        ProducerPerformance.ConfigPostProcessor configs = new 
ProducerPerformance.ConfigPostProcessor(parser, args);
+        assertEquals(ProducerPerformance.KeyDistribution.NONE, 
configs.keyDistribution);
+        assertNull(configs.recordKeyRange);
+    }
+
+    @Test
+    public void testKeyDistributionRange() throws IOException {
+        List<ProducerRecord<byte[], byte[]>> sentRecords = new ArrayList<>();
+        
doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));
+        doAnswer(invocation -> {
+            sentRecords.add(invocation.getArgument(0));
+            producerPerformanceSpy.cb.onCompletion(null, null);
+            return null;
+        }).when(producerMock).send(any(), any());
+
+        String[] args = new String[]{
+            "--topic", "Hello-Kafka",
+            "--num-records", "6",
+            "--throughput", "100",
+            "--record-size", "100",
+            "--key-distribution", "range",
+            "--record-key-range", "3",
+            "--bootstrap-server", "localhost:9000"};
+        producerPerformanceSpy.start(args);
+
+        assertEquals(6, sentRecords.size());
+        for (int i = 0; i < 6; i++) {
+            byte[] expectedKey = Integer.toString(i % 
3).getBytes(StandardCharsets.UTF_8);
+            assertArrayEquals(expectedKey, sentRecords.get(i).key(),
+                "Record " + i + " should have key " + (i % 3));
+        }
+    }
+
+    @Test
+    public void testKeyDistributionRandom() throws IOException {
+        List<ProducerRecord<byte[], byte[]>> sentRecords = new ArrayList<>();
+        
doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));
+        doAnswer(invocation -> {
+            sentRecords.add(invocation.getArgument(0));
+            producerPerformanceSpy.cb.onCompletion(null, null);
+            return null;
+        }).when(producerMock).send(any(), any());
+
+        String[] args = new String[]{
+            "--topic", "Hello-Kafka",
+            "--num-records", "10",
+            "--throughput", "100",
+            "--record-size", "100",
+            "--key-distribution", "random",
+            "--record-key-range", "5",
+            "--bootstrap-server", "localhost:9000"};
+        producerPerformanceSpy.start(args);
+
+        assertEquals(10, sentRecords.size());
+        for (ProducerRecord<byte[], byte[]> record : sentRecords) {
+            assertNotNull(record.key(), "Key should not be null for random 
distribution");
+            int keyValue = Integer.parseInt(new String(record.key(), 
StandardCharsets.UTF_8));
+            assertTrue(keyValue >= 0 && keyValue < 5, "Key should be in [0, 
5): " + keyValue);
+        }
+    }
+
+    @Test
+    public void testNullKeyWhenDistributionNone() throws IOException {
+        List<ProducerRecord<byte[], byte[]>> sentRecords = new ArrayList<>();
+        
doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));
+        doAnswer(invocation -> {
+            sentRecords.add(invocation.getArgument(0));
+            producerPerformanceSpy.cb.onCompletion(null, null);
+            return null;
+        }).when(producerMock).send(any(), any());
+
+        String[] args = new String[]{
+            "--topic", "Hello-Kafka",
+            "--num-records", "3",
+            "--throughput", "100",
+            "--record-size", "100",
+            "--bootstrap-server", "localhost:9000"};
+        producerPerformanceSpy.start(args);
+
+        assertEquals(3, sentRecords.size());
+        for (ProducerRecord<byte[], byte[]> record : sentRecords) {
+            assertNull(record.key(), "Key should be null when 
--key-distribution is 'none'");
+        }
+    }
+
+    @Test
+    public void testInvalidRecordKeyRange() {
+        ArgumentParser parser = ProducerPerformance.argParser();
+        String[] args = new String[]{
+            "--topic", "Hello-Kafka",
+            "--num-records", "5",
+            "--throughput", "100",
+            "--record-size", "100",
+            "--bootstrap-server", "localhost:9000",
+            "--key-distribution", "range",
+            "--record-key-range", "0"};
+        assertEquals("--record-key-range should be greater than zero.",
+            assertThrows(ArgumentParserException.class,
+                () -> new ProducerPerformance.ConfigPostProcessor(parser, 
args)).getMessage());
+    }
+
+    @Test
+    public void testKeyDistributionRequiresKeyRange() {
+        ArgumentParser parser = ProducerPerformance.argParser();
+        String[] args = new String[]{
+            "--topic", "Hello-Kafka",
+            "--num-records", "5",
+            "--throughput", "100",
+            "--record-size", "100",
+            "--bootstrap-server", "localhost:9000",
+            "--key-distribution", "range"};
+        assertEquals("--record-key-range is required when --key-distribution 
is 'range' or 'random'.",
+            assertThrows(ArgumentParserException.class,
+                () -> new ProducerPerformance.ConfigPostProcessor(parser, 
args)).getMessage());
+    }
+
+    @Test
+    public void testKeyRangeRequiresKeyDistribution() {
+        ArgumentParser parser = ProducerPerformance.argParser();
+        String[] args = new String[]{
+            "--topic", "Hello-Kafka",
+            "--num-records", "5",
+            "--throughput", "100",
+            "--record-size", "100",
+            "--bootstrap-server", "localhost:9000",
+            "--record-key-range", "5"};
+        assertEquals("--key-distribution must be 'range' or 'random' when 
--record-key-range is specified.",
+            assertThrows(ArgumentParserException.class,
+                () -> new ProducerPerformance.ConfigPostProcessor(parser, 
args)).getMessage());
+    }
+
+    @Test
+    public void testRandomSeedCustomValue() throws IOException, 
ArgumentParserException {
+        ArgumentParser parser = ProducerPerformance.argParser();
+        String[] args = new String[]{
+            "--topic", "Hello-Kafka",
+            "--num-records", "5",
+            "--throughput", "100",
+            "--record-size", "100",
+            "--bootstrap-server", "localhost:9000",
+            "--random-seed", "42"};
+        ProducerPerformance.ConfigPostProcessor configs = new 
ProducerPerformance.ConfigPostProcessor(parser, args);
+        assertEquals(42L, configs.randomSeed);
+    }
 }

Reply via email to