Ensure Kafka consumers are ready before starting the warmup traffic

merlimat · merlimat · commit 02e40c4aa3c4 · 2017-11-22T10:06:35.000-08:00
diff --git a/benchmark-framework/src/main/java/io/openmessaging/benchmark/WorkloadGenerator.java b/benchmark-framework/src/main/java/io/openmessaging/benchmark/WorkloadGenerator.java
@@ -32,7 +32,6 @@
 
 import org.HdrHistogram.Histogram;
 import org.HdrHistogram.Recorder;
-import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -78,6 +77,8 @@ public TestResult run() {
         List<BenchmarkConsumer> consumers = createConsumers(topics);
         List<BenchmarkProducer> producers = createProducers(topics);
 
+        ensureTopicsAreReady(producers, consumers);
+
         RateLimiter rateLimiter;
         AtomicBoolean runCompleted = new AtomicBoolean();
         if (workload.producerRate > 0) {
@@ -114,6 +115,27 @@ public TestResult run() {
         return result;
     }
 
+    private void ensureTopicsAreReady(List<BenchmarkProducer> producers, List<BenchmarkConsumer> consumers) {
+        log.info("Waiting for consumers to be ready");
+        // This is work around the fact that there's no way to have a consumer ready in Kafka without first publishing
+        // some message on the topic, which will then trigger the partitions assignement to the consumers
+
+        // In this case we just publish 1 message and then wait for consumers to receive the data
+        producers.forEach(producer -> producer.sendAsync(new byte[10]).thenRun(() -> totalMessagesSent.increment()));
+
+        long expectedMessages = workload.subscriptionsPerTopic * producers.size();
+
+        while (totalMessagesReceived.sum() < expectedMessages) {
+            try {
+                Thread.sleep(100);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        log.info("All consumers are ready");
+    }
+
     /**
      * Adjust the publish rate to a level that is sustainable, meaning that we can consume all the messages that are
      * being produced
@@ -246,6 +268,7 @@ private List<BenchmarkConsumer> createConsumers(List<String> topics) {
         }
 
         List<BenchmarkConsumer> consumers = futures.stream().map(CompletableFuture::join).collect(Collectors.toList());
+
         log.info("Created {} consumers in {} ms", consumers.size(), timer.elapsedMillis());
         return consumers;
     }
diff --git a/driver-kafka/kafka.yaml b/driver-kafka/kafka.yaml
@@ -25,11 +25,12 @@ driverClass: io.openmessaging.benchmark.driver.kafka.KafkaBenchmarkDriver
 
 replicationFactor: 1
 
-commonConfig: >
+commonConfig: |
     bootstrap.servers=localhost:9092
 
-producerConfig: >
+producerConfig: |
     acks=all
 
-consumerConfig: >
+consumerConfig: |
+    auto.offset.reset=earliest
     enable.auto.commit=false
