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);
+ }
}