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);
     }
 }
-
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..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
@@ -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.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 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);
 }