Added option to build & drain backlog in workload

merlimat · merlimat · commit 69241503a290 · 2017-11-22T12:09:26.000-08:00
diff --git a/benchmark-framework/src/main/java/io/openmessaging/benchmark/Workload.java b/benchmark-framework/src/main/java/io/openmessaging/benchmark/Workload.java
@@ -37,6 +37,15 @@ public class Workload {
     public int consumerPerSubscription;
 
     public int producerRate;
-    
+
+    /**
+     * If the consumer backlog is > 0, the generator will accumulate messages until the requested amount of storage is
+     * retained and then it will start the consumers to drain it.
+     * 
+     * The testDurationMinutes will be overruled to allow the test to complete when the consumer has drained all the
+     * backlog and it's on par with the producer
+     */
+    public long consumerBacklogSizeGB = 0;
+
     public int testDurationMinutes;
 }
diff --git a/benchmark-framework/src/main/java/io/openmessaging/benchmark/WorkloadGenerator.java b/benchmark-framework/src/main/java/io/openmessaging/benchmark/WorkloadGenerator.java
@@ -65,14 +65,19 @@ public class WorkloadGenerator implements ConsumerCallback, AutoCloseable {
     private final LongAdder totalMessagesReceived = new LongAdder();
 
     private boolean testCompleted = false;
+    private volatile boolean needToWaitForBacklogDraining = false;
 
     public WorkloadGenerator(String driverName, BenchmarkDriver benchmarkDriver, Workload workload) {
         this.driverName = driverName;
         this.benchmarkDriver = benchmarkDriver;
         this.workload = workload;
+
+        if (workload.consumerBacklogSizeGB > 0 && workload.producerRate == 0) {
+            throw new IllegalArgumentException("Cannot probe producer sustainable rate when building backlog");
+        }
     }
 
-    public TestResult run() {
+    public TestResult run() throws Exception {
         List<String> topics = createTopics();
         List<BenchmarkConsumer> consumers = createConsumers(topics);
         List<BenchmarkProducer> producers = createProducers(topics);
@@ -97,6 +102,17 @@ public TestResult run() {
         generateLoad(producers, TimeUnit.MINUTES.toSeconds(1), rateLimiter);
         runCompleted.set(true);
 
+        if (workload.consumerBacklogSizeGB > 0) {
+            log.info("Stopping all consumers to build backlog");
+            for (BenchmarkConsumer consumer : consumers) {
+                consumer.close();
+            }
+
+            executor.execute(() -> {
+                buildAndDrainBacklog(topics);
+            });
+        }
+
         log.info("----- Starting benchmark traffic ------");
         runCompleted.set(false);
 
@@ -288,6 +304,49 @@ private List<BenchmarkProducer> createProducers(List<String> topics) {
         return producers;
     }
 
+    private void buildAndDrainBacklog(List<String> topics) {
+        this.needToWaitForBacklogDraining = true;
+
+        long requestedBacklogSize = workload.consumerBacklogSizeGB * 1024 * 1024 * 1024;
+
+        while (true) {
+            long currentBacklogSize = (workload.subscriptionsPerTopic * totalMessagesSent.sum()
+                    - totalMessagesReceived.sum()) * workload.messageSize;
+
+            if (currentBacklogSize >= requestedBacklogSize) {
+                break;
+            }
+
+            try {
+                Thread.sleep(100);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        log.info("--- Start draining backlog ---");
+
+        createConsumers(topics);
+
+        final long minBacklog = 1000;
+
+        while (true) {
+            long currentBacklog = workload.subscriptionsPerTopic * totalMessagesSent.sum()
+                    - totalMessagesReceived.sum();
+            if (currentBacklog <= minBacklog) {
+                log.info("--- Completed backlog draining ---");
+                needToWaitForBacklogDraining = false;
+                return;
+            }
+
+            try {
+                Thread.sleep(100);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
     private TestResult generateLoad(List<BenchmarkProducer> producers, long testDurationsSeconds,
             RateLimiter rateLimiter) {
         Recorder recorder = new Recorder(TimeUnit.SECONDS.toMicros(30), 5);
@@ -386,7 +445,7 @@ private TestResult generateLoad(List<BenchmarkProducer> producers, long testDura
 
             reportHistogram.reset();
 
-            if (now >= testEndTime) {
+            if (now >= testEndTime && !needToWaitForBacklogDraining) {
                 testCompleted = true;
                 reportHistogram = cumulativeRecorder.getIntervalHistogram();
                 log.info(
diff --git a/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/KafkaBenchmarkConsumer.java b/driver-kafka/src/main/java/io/openmessaging/benchmark/driver/kafka/KafkaBenchmarkConsumer.java
@@ -22,6 +22,7 @@
 import java.util.Map;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
@@ -37,13 +38,15 @@ public class KafkaBenchmarkConsumer implements BenchmarkConsumer {
     private final KafkaConsumer<byte[], byte[]> consumer;
 
     private final ExecutorService executor;
+    private final Future<?> consumerTask;
+    private volatile boolean closing = false;
 
     public KafkaBenchmarkConsumer(KafkaConsumer<byte[], byte[]> consumer, ConsumerCallback callback) {
         this.consumer = consumer;
         this.executor = Executors.newSingleThreadExecutor();
 
-        this.executor.execute(() -> {
-            while (true) {
+        this.consumerTask = this.executor.submit(() -> {
+            while (!closing) {
                 ConsumerRecords<byte[], byte[]> records = consumer.poll(100);
 
                 Map<TopicPartition, OffsetAndMetadata> offsetMap = new HashMap<>();
@@ -63,8 +66,10 @@ public KafkaBenchmarkConsumer(KafkaConsumer<byte[], byte[]> consumer, ConsumerCa
 
     @Override
     public void close() throws Exception {
+        closing = true;
+        executor.shutdown();
+        consumerTask.get();
         consumer.close();
-        executor.shutdownNow();
     }
 
 }
diff --git a/workloads/simple-workload.yaml b/workloads/simple-workload.yaml
@@ -33,6 +33,8 @@ consumerPerSubscription : 1
 
 producerRate : 10000
 
+consumerBacklogSizeGB : 0
+
 testDurationMinutes : 5
 
 
