Allow setting arbitrary Kafka producer configs and callback

Author: SoerenHenning

load-generator/build.gradle

@@ -16,8 +16,8 @@ java {
     withSourcesJar()
 }
 
-group 'com.dynatrace'
-version '1.0-SNAPSHOT'
+group = 'com.dynatrace'
+version = '1.0-SNAPSHOT'
 
 repositories {
     mavenCentral()

load-generator/src/main/java/com/dynatrace/research/shufflebench/KafkaLoadGenerator.java

@@ -2,6 +2,8 @@
 
 import com.dynatrace.hash4j.hashing.Hashing;
 import com.dynatrace.research.shufflebench.record.RandomRecordGenerator;
+import com.dynatrace.research.shufflebench.record.Record;
+import io.smallrye.config.SmallRyeConfig;
 import org.eclipse.microprofile.config.Config;
 import org.eclipse.microprofile.config.ConfigProvider;
 import org.slf4j.Logger;
@@ -10,14 +12,14 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.*;
+import java.util.function.Consumer;
 
 public class KafkaLoadGenerator {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(KafkaLoadGenerator.class);
 
-  private final int executionTimeMs;
-
   private final String seedString;
 
   private final String kafkaBootstrapServers;
@@ -29,49 +31,82 @@ public class KafkaLoadGenerator {
 
   private final int recordSizeInBytes;
 
+  private final Map<String, String> kafkaProducerConfig;
+
   private final ScheduledExecutorService executor;
 
+  private final Consumer<Record> callback; // May be null
+
   private final List<KafkaSender> openKafkaSenders = new ArrayList<>();
   private final List<RecordSource> openRecordSources = new ArrayList<>();
 
+
   public KafkaLoadGenerator(
-          int executionTimeMs,
           String seedString,
           String kafkaBootstrapServers,
           String kafkaTopic,
           int numSources,
           int recordsPerSecondAndSource,
           int recordSizeInBytes,
-          int threadPoolSize
+          int threadPoolSize,
+          Map<String, String> kafkaProducerConfig
+  ) {
+    this(
+            seedString,
+            kafkaBootstrapServers,
+            kafkaTopic,
+            numSources,
+            recordsPerSecondAndSource,
+            recordSizeInBytes,
+            threadPoolSize,
+            kafkaProducerConfig,
+            null);
+  }
+
+  public KafkaLoadGenerator(
+          String seedString,
+          String kafkaBootstrapServers,
+          String kafkaTopic,
+          int numSources,
+          int recordsPerSecondAndSource,
+          int recordSizeInBytes,
+          int threadPoolSize,
+          Map<String, String> kafkaProducerConfig,
+          Consumer<Record> callback
   ) {
-    this.executionTimeMs = executionTimeMs;
     this.seedString = seedString;
     this.kafkaBootstrapServers = kafkaBootstrapServers;
     this.kafkaTopic = kafkaTopic;
     this.numSources = numSources;
     this.recordsPerSecondAndSource = recordsPerSecondAndSource;
     this.recordSizeInBytes = recordSizeInBytes;
     this.executor = new ScheduledThreadPoolExecutor(threadPoolSize);
+    this.kafkaProducerConfig = kafkaProducerConfig;
+    this.callback = callback;
   }
 
-  public void startBlocking() throws InterruptedException {
-    final KafkaSender kafkaSender = new KafkaSender(this.kafkaBootstrapServers, this.kafkaTopic);
+  public void startAsync() {
+    final KafkaSender kafkaSender = new KafkaSender(this.kafkaBootstrapServers, this.kafkaTopic, this.kafkaProducerConfig, this.callback);
     for (int sourceId = 0; sourceId < numSources; sourceId++) {
       final long seed = Hashing.komihash4_3().hashStream().putString(seedString).putInt(sourceId).getAsLong();
       final RecordSource recordSource = new RecordSource(
-          executor,
-          this.recordsPerSecondAndSource,
-          kafkaSender,
-          // new StaticRecordGenerator(),
-          new RandomRecordGenerator(seed, recordSizeInBytes),
-          "source" + sourceId);
+              executor,
+              this.recordsPerSecondAndSource,
+              kafkaSender,
+              // new StaticRecordGenerator(),
+              new RandomRecordGenerator(seed, recordSizeInBytes),
+              "source" + sourceId);
       openRecordSources.add(recordSource);
     }
     openKafkaSenders.add(kafkaSender);
-    Thread.sleep(this.executionTimeMs);
   }
 
-  public void stop() throws InterruptedException, IOException {
+  public void startBlocking(int executionTimeMs) throws InterruptedException {
+    this.startAsync();
+    Thread.sleep(executionTimeMs);
+  }
+
+  public void stop() {
     for (RecordSource recordSource : openRecordSources) {
       recordSource.stop();
     }
@@ -81,30 +116,31 @@ public void stop() throws InterruptedException, IOException {
     }
     openKafkaSenders.clear();
     executor.shutdown();
-    executor.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS);
+    try {
+      executor.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS);
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    }
   }
 
-  public static void main(String[] args) throws InterruptedException, IOException {
-    final Config config = ConfigProvider.getConfig();
+  public static void main(String[] args) throws InterruptedException {
+    final SmallRyeConfig config = ConfigProvider.getConfig().unwrap(SmallRyeConfig.class);
     final KafkaLoadGenerator kafkaLoadGenerator = new KafkaLoadGenerator(
-            config.getValue("execution.time.ms", Integer.class),
             config.getValue("seed.string", String.class),
             config.getValue("kafka.bootstrap.servers", String.class),
             config.getValue("kafka.topic", String.class),
             config.getValue("num.sources", Integer.class),
             config.getValue("num.records.per.source.second", Integer.class),
             config.getValue("record.size.bytes", Integer.class),
-            config.getValue("thread.pool.size", Integer.class)
+            config.getValue("thread.pool.size", Integer.class),
+            config.getOptionalValues("kafka.producer", String.class, String.class).orElse(Map.of())
     );
+    int executionTimeMs = config.getValue("execution.time.ms", Integer.class);
     Runtime.getRuntime().addShutdownHook(new Thread(() -> {
       LOGGER.info("Shut down load generator.");
-      try {
-        kafkaLoadGenerator.stop();
-      } catch (InterruptedException | IOException e) {
-        throw new RuntimeException(e);
-      }
+      kafkaLoadGenerator.stop();
     }));
-    kafkaLoadGenerator.startBlocking();
+    kafkaLoadGenerator.startBlocking(executionTimeMs);
     kafkaLoadGenerator.stop();
   }
 }

load-generator/src/main/java/com/dynatrace/research/shufflebench/KafkaSender.java

@@ -2,14 +2,12 @@
 
 import com.dynatrace.research.shufflebench.record.Record;
 import com.dynatrace.research.shufflebench.record.RecordSerializer;
-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.clients.producer.RecordMetadata;
+import org.apache.kafka.clients.producer.*;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
 import java.io.Closeable;
+import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.Future;
 import java.util.function.Consumer;
@@ -22,18 +20,34 @@ public class KafkaSender implements Consumer<Record>, Closeable {
 
   private final String topic;
 
-  public KafkaSender(final String bootstrapServers, final String topic) {
+  private final Consumer<Record> callback; // May be null
+
+  public KafkaSender(final String bootstrapServers, Map<String, String> kafkaProducerConfig, final String topic) {
+    this(bootstrapServers, topic, kafkaProducerConfig, null);
+  }
+
+  public KafkaSender(final String bootstrapServers, final String topic, Map<String, String> kafkaProducerConfig, final Consumer<Record> callback) {
     Properties properties = new Properties();
     properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
+    properties.putAll(kafkaProducerConfig);
     this.kafkaProducer = new KafkaProducer<>(properties, new NullKeySerializer(), new RecordSerializer());
     this.topic = topic;
+    this.callback = callback;
   }
 
   @Override
   public void accept(Record record) {
     ProducerRecord<Void, Record> producerRecord = new ProducerRecord<>(this.topic, record);
     LOGGER.debug("Send record to Kafka: {}", producerRecord);
-    Future<RecordMetadata> result = this.kafkaProducer.send(producerRecord);
+    Future<RecordMetadata> result = this.produceKafkaRecord(producerRecord);
+  }
+
+  private Future<RecordMetadata> produceKafkaRecord(ProducerRecord<Void, Record> producerRecord) {
+    if (callback == null) {
+      return this.kafkaProducer.send(producerRecord);
+    } else {
+      return this.kafkaProducer.send(producerRecord, (metadata, exception) -> callback.accept(producerRecord.value()));
+    }
   }
 
   @Override

load-generator/src/main/resources/META-INF/microprofile-config.properties

@@ -2,6 +2,9 @@ execution.time.ms=604800000
 seed.string=shuffle-load-generator-0
 kafka.bootstrap.servers=localhost:9092
 kafka.topic=input
+# Set any Kafka producer property with: kafka.producer."<...>"=<...>
+#kafka.producer."linger.ms"=100
+# To set via environment variable: KAFKA_PRODUCER__LINGER_MS__=100
 num.sources=1
 num.records.per.source.second=1000
 record.size.bytes=1024