From 280e98bf67258f7344270379aaaeceedd8b7c54e Mon Sep 17 00:00:00 2001 From: Masahiro Mori Date: Mon, 13 Apr 2026 15:44:16 +0900 Subject: [PATCH 1/3] feat: remove topics after benchmark --- .../benchmark/driver/kafka/Config.java | 2 + .../driver/kafka/KafkaBenchmarkDriver.java | 45 ++++++++++++++++++- 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/Config.java b/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/Config.java index d2fafafa3..757966e7a 100644 --- a/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/Config.java +++ b/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/Config.java @@ -16,6 +16,8 @@ public class Config { public short replicationFactor; + public boolean deleteTopicsOnClose = false; + public String topicConfig; public String commonConfig; diff --git a/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/KafkaBenchmarkDriver.java b/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/KafkaBenchmarkDriver.java index d506b5b37..0851e36bc 100644 --- a/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/KafkaBenchmarkDriver.java +++ b/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/KafkaBenchmarkDriver.java @@ -28,20 +28,27 @@ 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.Properties; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.Set; import org.apache.bookkeeper.stats.StatsLogger; import org.apache.kafka.clients.admin.AdminClient; +import org.apache.kafka.clients.admin.DeleteTopicsResult; import org.apache.kafka.clients.consumer.ConsumerConfig; import org.apache.kafka.clients.consumer.KafkaConsumer; import org.apache.kafka.clients.producer.KafkaProducer; import org.apache.kafka.clients.producer.ProducerConfig; +import org.apache.kafka.common.errors.UnknownTopicOrPartitionException; import org.apache.kafka.common.serialization.ByteArrayDeserializer; import org.apache.kafka.common.serialization.ByteArraySerializer; import org.apache.kafka.common.serialization.StringDeserializer; import org.apache.kafka.common.serialization.StringSerializer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class KafkaBenchmarkDriver implements BenchmarkDriver { @@ -52,6 +59,7 @@ public class KafkaBenchmarkDriver implements BenchmarkDriver { private List producers = Collections.synchronizedList(new ArrayList<>()); private List consumers = Collections.synchronizedList(new ArrayList<>()); + private final Set createdTopics = Collections.synchronizedSet(new HashSet<>()); private Properties topicProperties; private Properties producerProperties; @@ -111,7 +119,9 @@ public CompletableFuture createTopics(List topicInfos) { Map topicConfigs = new HashMap<>((Map) topicProperties); KafkaTopicCreator topicCreator = new KafkaTopicCreator(admin, topicConfigs, config.replicationFactor); - return topicCreator.create(topicInfos); + return topicCreator + .create(topicInfos) + .thenRun(() -> topicInfos.forEach(topicInfo -> createdTopics.add(topicInfo.getTopic()))); } @Override @@ -158,9 +168,40 @@ public void close() throws Exception { for (BenchmarkConsumer consumer : consumers) { consumer.close(); } + if (config != null && config.deleteTopicsOnClose) { + deleteCreatedTopics(); + } admin.close(); } + private void deleteCreatedTopics() { + if (createdTopics.isEmpty()) { + return; + } + + final Set topicsToDelete; + synchronized (createdTopics) { + topicsToDelete = new HashSet<>(createdTopics); + } + + try { + DeleteTopicsResult deleteTopicsResult = admin.deleteTopics(topicsToDelete); + deleteTopicsResult.all().get(); + log.info("Deleted {} benchmark topics", topicsToDelete.size()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + log.warn("Interrupted while deleting benchmark topics {}", topicsToDelete, e); + } catch (ExecutionException e) { + if (e.getCause() instanceof UnknownTopicOrPartitionException) { + log.warn("Some benchmark topics were already deleted: {}", topicsToDelete); + } else { + log.warn("Failed deleting benchmark topics {}", topicsToDelete, e); + } + } catch (Exception e) { + log.warn("Failed deleting benchmark topics {}", topicsToDelete, e); + } + } + private static String applyZoneId(String clientId, String zoneId) { return clientId.replace(ZONE_ID_TEMPLATE, zoneId); } @@ -168,4 +209,6 @@ private static String applyZoneId(String clientId, String zoneId) { private static final ObjectMapper mapper = new ObjectMapper(new YAMLFactory()) .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false); + + private static final Logger log = LoggerFactory.getLogger(KafkaBenchmarkDriver.class); } From 0d0a7715e14f4edffbeeb185d911f2355e89a525 Mon Sep 17 00:00:00 2001 From: Masahiro Mori Date: Tue, 14 Apr 2026 11:41:31 +0900 Subject: [PATCH 2/3] fix: sort imports --- .../benchmark/driver/kafka/KafkaBenchmarkDriver.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/KafkaBenchmarkDriver.java b/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/KafkaBenchmarkDriver.java index 0851e36bc..dfa6ecef1 100644 --- a/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/KafkaBenchmarkDriver.java +++ b/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/KafkaBenchmarkDriver.java @@ -32,9 +32,9 @@ import java.util.List; import java.util.Map; import java.util.Properties; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; -import java.util.Set; import org.apache.bookkeeper.stats.StatsLogger; import org.apache.kafka.clients.admin.AdminClient; import org.apache.kafka.clients.admin.DeleteTopicsResult; From e0bfc66535fe8a1782cff14bb9f02a41cb665204 Mon Sep 17 00:00:00 2001 From: Masahiro Mori Date: Wed, 24 Jun 2026 08:08:20 +0900 Subject: [PATCH 3/3] fix(DistributedWorkersEnsemble): stop the leader worker after stopping other workers --- .../io/openmessaging/benchmark/Workload.java | 7 ++-- .../payload/MessageSizeDistribution.java | 15 ++++--- .../worker/DistributedWorkersEnsemble.java | 42 ++++++++++++++++++- .../benchmark/worker/LocalWorker.java | 4 +- .../commands/ProducerWorkAssignment.java | 6 +-- .../payload/MessageSizeDistributionTest.java | 1 - 6 files changed, 55 insertions(+), 20 deletions(-) diff --git a/benchmark-framework/src/main/java/io/openmessaging/benchmark/Workload.java b/benchmark-framework/src/main/java/io/openmessaging/benchmark/Workload.java index 7b877a197..3562ffb53 100644 --- a/benchmark-framework/src/main/java/io/openmessaging/benchmark/Workload.java +++ b/benchmark-framework/src/main/java/io/openmessaging/benchmark/Workload.java @@ -31,10 +31,9 @@ public class Workload { public int messageSize; /** - * Message size distribution for variable-sized payloads. - * Keys are size ranges (e.g., "0-256", "256-1024", "1KB-4KB"), - * values are relative weights. - * Mutually exclusive with messageSize - if set, messageSize is ignored. + * Message size distribution for variable-sized payloads. Keys are size ranges (e.g., "0-256", + * "256-1024", "1KB-4KB"), values are relative weights. Mutually exclusive with messageSize - if + * set, messageSize is ignored. */ public Map messageSizeDistribution; diff --git a/benchmark-framework/src/main/java/io/openmessaging/benchmark/utils/payload/MessageSizeDistribution.java b/benchmark-framework/src/main/java/io/openmessaging/benchmark/utils/payload/MessageSizeDistribution.java index 56355e4b6..6569d92c1 100644 --- a/benchmark-framework/src/main/java/io/openmessaging/benchmark/utils/payload/MessageSizeDistribution.java +++ b/benchmark-framework/src/main/java/io/openmessaging/benchmark/utils/payload/MessageSizeDistribution.java @@ -13,15 +13,17 @@ */ package io.openmessaging.benchmark.utils.payload; + import java.util.ArrayList; import java.util.List; import java.util.Map; /** - * Parses and represents a message size distribution from workload config. - * Creates one payload size per bucket and provides weights for runtime selection. + * Parses and represents a message size distribution from workload config. Creates one payload size + * per bucket and provides weights for runtime selection. * *

Example configuration: + * * 

  * messageSizeDistribution:
  *   "0-256": 234
@@ -34,9 +36,7 @@ public class MessageSizeDistribution {
     private final List buckets;
     private final int totalWeight;
 
-    /**
-     * Represents a single size bucket with min/max range and weight.
-     */
+    /** Represents a single size bucket with min/max range and weight. */
     public static class Bucket {
         public final int minSize;
         public final int maxSize;
@@ -147,8 +147,8 @@ public List getBucketSizes() {
     }
 
     /**
-     * Returns list of max sizes, one per bucket (for payload generation).
-     * Using max sizes ensures the system is tested with the largest messages in each bucket range.
+     * Returns list of max sizes, one per bucket (for payload generation). Using max sizes ensures the
+     * system is tested with the largest messages in each bucket range.
      *
      * @return list of max sizes per bucket
      */
@@ -217,4 +217,3 @@ public List getBuckets() {
         return buckets;
     }
 }
-
diff --git a/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/DistributedWorkersEnsemble.java b/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/DistributedWorkersEnsemble.java
index c41a0da51..f326ccc8c 100644
--- a/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/DistributedWorkersEnsemble.java
+++ b/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/DistributedWorkersEnsemble.java
@@ -17,6 +17,9 @@
 import static java.util.stream.Collectors.joining;
 
 import com.beust.jcommander.internal.Maps;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Lists;
@@ -36,6 +39,8 @@
 import org.slf4j.LoggerFactory;
 
 public class DistributedWorkersEnsemble implements Worker {
+    private static final String KAFKA_BENCHMARK_DRIVER_CLASS =
+            "io.openmessaging.benchmark.driver.kafka.KafkaBenchmarkDriver";
     private final Thread shutdownHook = new Thread(this::stopAll);
     private final List workers;
     private final List producerWorkers;
@@ -43,6 +48,7 @@ public class DistributedWorkersEnsemble implements Worker {
     private final Worker leader;
 
     private int numberOfUsedProducerWorkers;
+    private volatile boolean stopLeaderLast;
 
     public DistributedWorkersEnsemble(List workers, boolean extraConsumerWorkers) {
         Preconditions.checkArgument(workers.size() > 1);
@@ -75,6 +81,7 @@ static int getNumberOfProducerWorkers(List workers, boolean extraConsume
 
     @Override
     public void initializeDriver(File configurationFile) throws IOException {
+        stopLeaderLast = shouldStopLeaderLast(configurationFile);
         workers.parallelStream()
                 .forEach(
                         w -> {
@@ -165,7 +172,39 @@ public void adjustPublishRate(double publishRate) throws IOException {
 
     @Override
     public void stopAll() {
-        workers.parallelStream().forEach(Worker::stopAll);
+        if (!stopLeaderLast) {
+            workers.parallelStream().forEach(Worker::stopAll);
+            return;
+        }
+
+        RuntimeException stopError = null;
+
+        try {
+            // The leader owns topic lifecycle, so stop it after every other worker is done.
+            workers.parallelStream().filter(worker -> worker != leader).forEach(Worker::stopAll);
+        } catch (RuntimeException e) {
+            stopError = e;
+        }
+
+        try {
+            leader.stopAll();
+        } catch (RuntimeException e) {
+            if (stopError != null) {
+                stopError.addSuppressed(e);
+            } else {
+                stopError = e;
+            }
+        }
+
+        if (stopError != null) {
+            throw stopError;
+        }
+    }
+
+    static boolean shouldStopLeaderLast(File configurationFile) throws IOException {
+        JsonNode configuration = mapper.readTree(configurationFile);
+        return KAFKA_BENCHMARK_DRIVER_CLASS.equals(configuration.path("driverClass").asText())
+                && configuration.path("deleteTopicsOnClose").asBoolean(false);
     }
 
     @Override
@@ -290,4 +329,5 @@ public void close() throws Exception {
     }
 
     private static final Logger log = LoggerFactory.getLogger(DistributedWorkersEnsemble.class);
+    private static final ObjectMapper mapper = new ObjectMapper(new YAMLFactory());
 }
diff --git a/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/LocalWorker.java b/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/LocalWorker.java
index af64369c4..70e91f368 100644
--- a/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/LocalWorker.java
+++ b/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/LocalWorker.java
@@ -240,9 +240,7 @@ private void submitProducersToExecutor(
                                             idx = r.nextInt(payloadCount);
                                         }
                                         messageProducer.sendMessage(
-                                                p,
-                                                Optional.ofNullable(keyDistributor.next()),
-                                                payloads.get(idx));
+                                                p, Optional.ofNullable(keyDistributor.next()), payloads.get(idx));
                                     });
                         }
                     } catch (Throwable t) {
diff --git a/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/commands/ProducerWorkAssignment.java b/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/commands/ProducerWorkAssignment.java
index 319bd7cc8..38e13705b 100644
--- a/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/commands/ProducerWorkAssignment.java
+++ b/benchmark-framework/src/main/java/io/openmessaging/benchmark/worker/commands/ProducerWorkAssignment.java
@@ -22,9 +22,9 @@ public class ProducerWorkAssignment {
     public List payloadData;
 
     /**
-     * Weights for weighted payload selection. If null, uniform random selection is used.
-     * Each weight corresponds to the payload at the same index in payloadData.
-     * Used for message size distribution feature.
+     * Weights for weighted payload selection. If null, uniform random selection is used. Each weight
+     * corresponds to the payload at the same index in payloadData. Used for message size distribution
+     * feature.
      */
     public int[] payloadWeights;
 
diff --git a/benchmark-framework/src/test/java/io/openmessaging/benchmark/utils/payload/MessageSizeDistributionTest.java b/benchmark-framework/src/test/java/io/openmessaging/benchmark/utils/payload/MessageSizeDistributionTest.java
index affa77536..8c72a4f00 100644
--- a/benchmark-framework/src/test/java/io/openmessaging/benchmark/utils/payload/MessageSizeDistributionTest.java
+++ b/benchmark-framework/src/test/java/io/openmessaging/benchmark/utils/payload/MessageSizeDistributionTest.java
@@ -219,4 +219,3 @@ void weightedSelectionProducesCorrectDistribution() {
         assertThat(ratio1).isBetween(0.49, 0.51);
     }
 }
-