Merge pull request #288 from rabbitmq/cmessages

Add --cmessages to limit the number of consumed messages

Author: acogoluegnes

ci/start-broker.sh

@@ -14,6 +14,7 @@ wait_for_message() {
 
 make -C "${PWD}"/tls-gen/basic
 
+rm -rf rabbitmq-configuration
 mkdir -p rabbitmq-configuration/tls
 cp -R "${PWD}"/tls-gen/basic/result/* rabbitmq-configuration/tls
 chmod o+r rabbitmq-configuration/tls/*

src/main/java/com/rabbitmq/stream/perf/StreamPerfTest.java

@@ -705,6 +705,13 @@ void setTcpNoDelay(String input) throws Exception {
       converter = Converters.GreaterThanOrEqualToZeroIntegerTypeConverter.class)
   private long pmessages;
 
+  @CommandLine.Option(
+      names = {"--cmessages", "-D"},
+      description = "consumer message count, default is 0 (no limit)",
+      defaultValue = "0",
+      converter = Converters.GreaterThanOrEqualToZeroIntegerTypeConverter.class)
+  private long cmessages;
+
   private MetricsCollector metricsCollector;
   private PerformanceMetrics performanceMetrics;
   private List<Monitoring> monitorings;
@@ -1094,10 +1101,12 @@ public Integer call() throws Exception {
 
       CompletionHandler completionHandler;
       ConcurrentMap<String, Integer> completionReasons = new ConcurrentHashMap<>();
-      if (isRunTimeLimited() || this.pmessages > 0) {
+      if (isRunTimeLimited() || this.pmessages > 0 || this.cmessages > 0) {
+        int countLimit = this.pmessages > 0 ? this.producers : 0;
+        countLimit += this.cmessages > 0 ? this.consumers : 0;
         completionHandler =
             new CompletionHandler.DefaultCompletionHandler(
-                this.time, this.producers, completionReasons);
+                this.time, countLimit, completionReasons);
       } else {
         completionHandler = new CompletionHandler.NoLimitCompletionHandler(completionReasons);
       }
@@ -1319,6 +1328,20 @@ public Integer call() throws Exception {
                                   .builder();
                         }
 
+                        java.util.function.Consumer<MessageHandler.Context> messageReceivedCallback;
+                        if (this.cmessages > 0) {
+                          AtomicLong messageCount = new AtomicLong(0);
+                          messageReceivedCallback =
+                              ctx -> {
+                                if (messageCount.incrementAndGet() == this.cmessages) {
+                                  completionHandler.countDown("Consumer reached message limit");
+                                  ctx.consumer().close();
+                                }
+                              };
+                        } else {
+                          messageReceivedCallback = ctx -> {};
+                        }
+
                         Runnable latencyWorker = Utils.latencyWorker(this.consumerLatency);
                         consumerBuilder =
                             consumerBuilder.messageHandler(
@@ -1333,6 +1356,7 @@ public Integer call() throws Exception {
                                     // tool
                                   }
                                   metrics.offset(context.offset());
+                                  messageReceivedCallback.accept(context);
                                   latencyWorker.run();
                                 });
 

src/test/java/com/rabbitmq/stream/perf/StreamPerfTestTest.java

@@ -573,6 +573,20 @@ void shouldPublishExpectedNumberOfMessagesIfOptionIsSet(long pmessages, int prod
     }
   }
 
+  @ParameterizedTest
+  @CsvSource({"100, 1", "100, 2"})
+  void shouldConsumeExpectedNumberOfMessagesIfOptionIsSet(long cmessages, int consumerCount)
+      throws Exception {
+    run(
+        builder()
+            .cmessages(cmessages)
+            .consumers(consumerCount)
+            .producers(1)
+            .pmessages(cmessages * 2));
+    waitUntilStreamExists(s);
+    waitRunEnds();
+  }
+
   private static <T> Consumer<T> wrap(CallableConsumer<T> action) {
     return t -> {
       try {
@@ -820,6 +834,11 @@ ArgumentsBuilder pmessages(long messages) {
       return this;
     }
 
+    ArgumentsBuilder cmessages(long messages) {
+      arguments.put("cmessages", String.valueOf(messages));
+      return this;
+    }
+
     String build() {
       return this.arguments.entrySet().stream()
           .map(e -> "--" + e.getKey() + (e.getValue().isEmpty() ? "" : (" " + e.getValue())))