Support filtering in performance tool

acogoluegnes · acogoluegnes · commit dedf25c590d2 · 2023-06-20T10:19:21.000+02:00
With --filter-value-set and --filter-values flags.
diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java
@@ -151,7 +151,9 @@ StreamConsumerBuilder lazyInit(boolean lazyInit) {
 
   @Override
   public FilterConfiguration filter() {
-    this.filterConfiguration = new DefaultFilterConfiguration(this);
+    if (this.filterConfiguration == null) {
+      this.filterConfiguration = new DefaultFilterConfiguration(this);
+    }
     return this.filterConfiguration;
   }
 
diff --git a/src/main/java/com/rabbitmq/stream/perf/StreamPerfTest.java b/src/main/java/com/rabbitmq/stream/perf/StreamPerfTest.java
@@ -439,6 +439,18 @@ public class StreamPerfTest implements Callable<Integer> {
   @ArgGroup(exclusive = false, multiplicity = "0..1")
   InstanceSyncOptions instanceSyncOptions;
 
+  @CommandLine.Option(
+      names = {"--filter-value-set", "-fvs"},
+      description = "filter value set for publishers, range (e.g. 1..15) are accepted",
+      converter = Utils.FilterValueSetConverter.class)
+  private List<String> filterValueSet;
+
+  @CommandLine.Option(
+      names = {"--filter-values", "-fv"},
+      description = "filter values for consumers",
+      split = ",")
+  private List<String> filterValues;
+
   static class InstanceSyncOptions {
 
     @CommandLine.Option(
@@ -589,7 +601,6 @@ public Integer call() throws Exception {
     maybeDisplayVersion();
     maybeDisplayEnvironmentVariablesHelp();
     overridePropertiesWithEnvironmentVariables();
-
     Codec codec = createCodec(this.codecClass);
 
     ByteBufAllocator byteBufAllocator = ByteBufAllocator.DEFAULT;
@@ -876,19 +887,43 @@ public Integer call() throws Exception {
                           producerBuilder.name(producerName).confirmTimeout(Duration.ZERO);
                     }
 
-                    java.util.function.Consumer<MessageBuilder> messageBuilderConsumer;
+                    java.util.function.Consumer<MessageBuilder> messageBuilderConsumerTemp;
                     if (this.superStreams) {
                       producerBuilder
                           .superStream(stream)
                           .routing(msg -> msg.getProperties().getMessageIdAsString());
                       AtomicLong messageIdSequence = new AtomicLong(0);
-                      messageBuilderConsumer =
+                      messageBuilderConsumerTemp =
                           mg -> mg.properties().messageId(messageIdSequence.getAndIncrement());
                     } else {
-                      messageBuilderConsumer = mg -> {};
+                      messageBuilderConsumerTemp = mg -> {};
                       producerBuilder.stream(stream);
                     }
 
+                    if (this.filterValueSet != null && this.filterValueSet.size() > 0) {
+                      producerBuilder =
+                          producerBuilder.filterValue(msg -> msg.getProperties().getTo());
+                      List<String> values = new ArrayList<>(this.filterValueSet);
+                      AtomicInteger count = new AtomicInteger();
+                      int subSetSize = Utils.filteringSubSetSize(values.size());
+                      int messageCountCycle = Utils.filteringPublishingCycle(this.rate);
+                      List<String> subSet = new ArrayList<>(subSetSize);
+                      java.util.function.Consumer<MessageBuilder> filteringMessageBuilderConsumer =
+                          b -> {
+                            if (Integer.remainderUnsigned(
+                                    count.getAndIncrement(), messageCountCycle)
+                                == 0) {
+                              Collections.shuffle(values);
+                              subSet.clear();
+                              subSet.addAll(values.subList(0, subSetSize));
+                            }
+                            b.properties()
+                                .to(subSet.get(Integer.remainderUnsigned(count.get(), subSetSize)));
+                          };
+                      messageBuilderConsumerTemp =
+                          messageBuilderConsumerTemp.andThen(filteringMessageBuilderConsumer);
+                    }
+
                     Producer producer =
                         producerBuilder
                             .subEntrySize(this.subEntrySize)
@@ -898,9 +933,9 @@ public Integer call() throws Exception {
                             .maxUnconfirmedMessages(this.confirms)
                             .build();
 
-                    AtomicLong messageCount = new AtomicLong(0);
                     ConfirmationHandler confirmationHandler;
                     if (this.confirmLatency) {
+                      AtomicLong messageCount = new AtomicLong(0);
                       final PerformanceMetrics metrics = this.performanceMetrics;
                       final int divisor = Utils.downSamplingDivisor(this.rate);
                       confirmationHandler =
@@ -936,6 +971,8 @@ public Integer call() throws Exception {
 
                     producers.add(producer);
 
+                    java.util.function.Consumer<MessageBuilder> messageBuilderConsumer =
+                        messageBuilderConsumerTemp;
                     return (Runnable)
                         () -> {
                           final int msgSize = this.messageSize;
@@ -1046,6 +1083,8 @@ public Integer call() throws Exception {
                                   }
                                 });
 
+                        consumerBuilder = maybeConfigureForFiltering(consumerBuilder);
+
                         Consumer consumer = consumerBuilder.build();
                         return consumer;
                       })
@@ -1122,6 +1161,37 @@ public Integer call() throws Exception {
     return 0;
   }
 
+  private ConsumerBuilder maybeConfigureForFiltering(ConsumerBuilder consumerBuilder) {
+    if (this.filterValues != null && this.filterValues.size() > 0) {
+      consumerBuilder =
+          consumerBuilder.filter().values(this.filterValues.toArray(new String[0])).builder();
+
+      if (this.filterValues.size() == 1) {
+        String filterValue = filterValues.get(0);
+        consumerBuilder =
+            consumerBuilder
+                .filter()
+                .postFilter(msg -> filterValue.equals(msg.getProperties().getTo()))
+                .builder();
+      } else {
+        consumerBuilder =
+            consumerBuilder
+                .filter()
+                .postFilter(
+                    msg -> {
+                      for (String filterValue : this.filterValues) {
+                        if (filterValue.equals(msg.getProperties().getTo())) {
+                          return true;
+                        }
+                      }
+                      return false;
+                    })
+                .builder();
+      }
+    }
+    return consumerBuilder;
+  }
+
   private void createStream(Environment environment, String stream) {
     StreamCreator streamCreator =
         environment.streamCreator().stream(stream)
diff --git a/src/main/java/com/rabbitmq/stream/perf/Utils.java b/src/main/java/com/rabbitmq/stream/perf/Utils.java
@@ -473,6 +473,33 @@ public BiFunction<String, Integer, String> convert(String input) {
     }
   }
 
+  static class FilterValueSetConverter implements ITypeConverter<List<String>> {
+
+    @Override
+    public List<String> convert(String value) {
+      if (value == null || value.trim().isEmpty()) {
+        return Collections.emptyList();
+      }
+      if (value.contains("..")) {
+        String[] range = value.split("\\.\\.");
+        String errorMessage = "'" + value + "' is not valid, valid example values: 1..10, 1..20";
+        if (range.length != 2) {
+          throw new CommandLine.TypeConversionException(errorMessage);
+        }
+        int start, end;
+        try {
+          start = Integer.parseInt(range[0]);
+          end = Integer.parseInt(range[1]) + 1;
+          return IntStream.range(start, end).mapToObj(String::valueOf).collect(Collectors.toList());
+        } catch (NumberFormatException e) {
+          throw new CommandLine.TypeConversionException(errorMessage);
+        }
+      } else {
+        return Arrays.stream(value.split(",")).collect(Collectors.toList());
+      }
+    }
+  }
+
   static class SniServerNamesConverter implements ITypeConverter<List<SNIServerName>> {
 
     @Override
@@ -918,4 +945,24 @@ static InstanceSynchronization defaultInstanceSynchronization(
       throw new RuntimeException(e);
     }
   }
+
+  static int filteringPublishingCycle(int rate) {
+    if (rate == 0) {
+      return 100_000;
+    } else if (rate <= 10) {
+      return 10;
+    } else {
+      return rate / 10;
+    }
+  }
+
+  static int filteringSubSetSize(int setSize) {
+    if (setSize <= 3) {
+      return 1;
+    } else if (setSize > 10) {
+      return (int) (setSize * 0.70);
+    } else {
+      return setSize - 3;
+    }
+  }
 }
diff --git a/src/test/java/com/rabbitmq/stream/perf/StreamPerfTestTest.java b/src/test/java/com/rabbitmq/stream/perf/StreamPerfTestTest.java
@@ -511,6 +511,17 @@ void nativeEpollWorksOnLinux() throws Exception {
     assertThat(streamExists(s)).isTrue();
   }
 
+  @Test
+  @BrokerVersionAtLeast(BrokerVersion.RABBITMQ_3_13_0)
+  void shouldNotFailWhenFilteringIsActivated() throws Exception {
+    Future<?> run = run(builder().filterValueSet("1..15").filterValues("4"));
+    waitUntilStreamExists(s);
+    waitOneSecond();
+    run.cancel(true);
+    waitRunEnds();
+    assertThat(consoleOutput()).containsIgnoringCase("summary:");
+  }
+
   private static <T> Consumer<T> wrap(CallableConsumer<T> action) {
     return t -> {
       try {
@@ -728,6 +739,16 @@ ArgumentsBuilder nativeEpoll() {
       return this;
     }
 
+    ArgumentsBuilder filterValueSet(String... values) {
+      arguments.put("filter-value-set", String.join(",", values));
+      return this;
+    }
+
+    ArgumentsBuilder filterValues(String... values) {
+      arguments.put("filter-values", String.join(",", values));
+      return this;
+    }
+
     String build() {
       return this.arguments.entrySet().stream()
           .map(e -> "--" + e.getKey() + (e.getValue().isEmpty() ? "" : (" " + e.getValue())))
diff --git a/src/test/java/com/rabbitmq/stream/perf/UtilsTest.java b/src/test/java/com/rabbitmq/stream/perf/UtilsTest.java
@@ -385,6 +385,27 @@ void creditsConverterKo(String input) {
         .isInstanceOf(TypeConversionException.class);
   }
 
+  @ParameterizedTest
+  @CsvSource({"0,100000", "100,10", "1000,100", "50,5", "10,10", "11,1", "5,10"})
+  void filteringPublishingCycle(int rate, int expected) {
+    assertThat(Utils.filteringPublishingCycle(rate)).isEqualTo(expected);
+  }
+
+  @ParameterizedTest
+  @CsvSource({"3,1", "2,1", "4,1", "7,4", "10,7", "15,10"})
+  void filteringSubSetSize(int setSize, int expected) {
+    assertThat(Utils.filteringSubSetSize(setSize)).isEqualTo(expected);
+  }
+
+  @Test
+  void filterValueSetConverter() throws Exception {
+    CommandLine.ITypeConverter<List<String>> converter = new Utils.FilterValueSetConverter();
+    assertThat(converter.convert("one")).containsExactly("one");
+    assertThat(converter.convert("one,two,three")).containsExactly("one", "two", "three");
+    assertThat(converter.convert("1..10")).hasSize(10).contains("1", "2", "10");
+    assertThat(converter.convert("5..10")).hasSize(6).contains("5", "6", "10");
+  }
+
   @Command(name = "test-command")
   static class TestCommand {
 

Original file line number	Diff line number	Diff line change
`@@ -151,7 +151,9 @@ StreamConsumerBuilder lazyInit(boolean lazyInit) {`
`151`	`151`
`152`	`152`	`@Override`
`153`	`153`	`public FilterConfiguration filter() {`
`154`		`- this.filterConfiguration = new DefaultFilterConfiguration(this);`
	`154`	`+ if (this.filterConfiguration == null) {`
	`155`	`+ this.filterConfiguration = new DefaultFilterConfiguration(this);`
	`156`	`+ }`
`155`	`157`	`return this.filterConfiguration;`
`156`	`158`	`}`
`157`	`159`