feat: add option for datagen to produce indefinitely (MINOR) (#3307)

BREAKING CHANGE: If no value is passed for the KSQL datagen option `iterations`, datagen will now produce indefinitely, rather than terminating after a default of 1,000,000 rows.
confluentinc · Sep 6, 2019 · 6281738 · 6281738
1 parent 0d0b1c3
commit 6281738
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 23 deletions.
diff --git a/ksql-examples/src/main/java/io/confluent/ksql/datagen/DataGen.java b/ksql-examples/src/main/java/io/confluent/ksql/datagen/DataGen.java
@@ -149,7 +149,8 @@ private static void usage() {
             + "'delimited') " + newLine
         + "topic=<kafka topic name> " + newLine
         + "key=<name of key column> " + newLine
-        + "[iterations=<number of rows> (defaults to 1,000,000)] " + newLine
+        + "[iterations=<number of rows> (if no value is specified, datagen will produce "
+            + "indefinitely)] " + newLine
         + "[maxInterval=<Max time in ms between rows> (defaults to 500)] " + newLine
         + "[propertiesFile=<file specifying Kafka client properties>] " + newLine
         + "[nThreads=<number of producer threads to start>] " + newLine
@@ -266,7 +267,7 @@ private Builder() {
         valueFormat = null;
         topicName = null;
         keyName = null;
-        iterations = 1000000;
+        iterations = -1;
         maxInterval = -1;
         schemaRegistryUrl = "http://localhost:8081";
         propertiesFile = null;

diff --git a/ksql-examples/src/main/java/io/confluent/ksql/datagen/DataGenProducer.java b/ksql-examples/src/main/java/io/confluent/ksql/datagen/DataGenProducer.java
@@ -79,33 +79,65 @@ public void populateTopic(
         valueSerializer
     );
 
-    for (int i = 0; i < messageCount; i++) {
-      rateLimiter.ifPresent(RateLimiter::acquire);
+    if (messageCount != -1) {
+      for (int i = 0; i < messageCount; i++) {
+        produceOne(
+            rowGenerator,
+            producer,
+            kafkaTopicName,
+            maxInterval,
+            printRows,
+            rateLimiter
+        );
+      }
+    } else {
+      while (true) {
+        produceOne(
+            rowGenerator,
+            producer,
+            kafkaTopicName,
+            maxInterval,
+            printRows,
+            rateLimiter
+        );
+      }
+    }
 
-      final Pair<Struct, GenericRow> genericRowPair = rowGenerator.generateRow();
+    producer.flush();
+    producer.close();
+  }
 
-      final ProducerRecord<Struct, GenericRow> producerRecord = new ProducerRecord<>(
-          kafkaTopicName,
-          genericRowPair.getLeft(),
-          genericRowPair.getRight()
-      );
+  private void produceOne(
+      final RowGenerator rowGenerator,
+      final KafkaProducer<Struct, GenericRow> producer,
+      final String kafkaTopicName,
+      final long maxInterval,
+      final boolean printRows,
+      final Optional<RateLimiter> rateLimiter
+  ) {
+    rateLimiter.ifPresent(RateLimiter::acquire);
 
-      producer.send(producerRecord,
-          new LoggingCallback(kafkaTopicName,
-              genericRowPair.getLeft(),
-              genericRowPair.getRight(),
-              printRows));
+    final Pair<Struct, GenericRow> genericRowPair = rowGenerator.generateRow();
 
-      try {
-        final long interval = maxInterval < 0 ? INTER_MESSAGE_MAX_INTERVAL : maxInterval;
+    final ProducerRecord<Struct, GenericRow> producerRecord = new ProducerRecord<>(
+        kafkaTopicName,
+        genericRowPair.getLeft(),
+        genericRowPair.getRight()
+    );
 
-        Thread.sleep((long) (interval * Math.random()));
-      } catch (final InterruptedException e) {
-        // Ignore the exception.
-      }
+    producer.send(producerRecord,
+        new LoggingCallback(kafkaTopicName,
+            genericRowPair.getLeft(),
+            genericRowPair.getRight(),
+            printRows));
+
+    try {
+      final long interval = maxInterval < 0 ? INTER_MESSAGE_MAX_INTERVAL : maxInterval;
+
+      Thread.sleep((long) (interval * Math.random()));
+    } catch (final InterruptedException e) {
+      // Ignore the exception.
     }
-    producer.flush();
-    producer.close();
   }
 
   private Serializer<Struct> getKeySerializer() {