Support environment variables in performance tool

acogoluegnes · acogoluegnes · commit ac7d4184981a · 2021-10-26T10:13:50.000+02:00
Environment variables can override command line arguments. Fixes #37
diff --git a/src/docs/asciidoc/performance-tool.adoc b/src/docs/asciidoc/performance-tool.adoc
@@ -477,6 +477,31 @@ This can be useful to inspect threads if the tool seems blocked.
 The endpoint can then be declared in a Prometheus instance to scrape the metrics.
 * `--monitoring-port`: set the port to use for the web server.
 
+===== Using Environment Variables as Options
+
+Environment variables can sometimes be easier to work with than command line options.
+This is especially true when using a manifest file for configuration (with Docker Compose or Kubernetes) and the number of options used grows.
+
+The performance tool automatically uses environment variables that match the snake case version of its long options.
+E.g. it automatically picks up the value of the `BATCH_SIZE` environment variable for the `--batch-size` option, but only if the environment variable is defined.
+
+You can list the environment variables that the tool picks up with the following command:
+
+----
+java -jar stream-perf-test.jar --environment-variables
+----
+
+The short version of the option is `-env`.
+
+To avoid collisions with environment variables that already exist, it is possible to specify a prefix for the environment variables that the tool looks up.
+This prefix is defined with the `RABBITMQ_STREAM_PERF_TEST_ENV_PREFIX` environment variable, e.g.:
+
+----
+RABBITMQ_STREAM_PERF_TEST_ENV_PREFIX="STREAM_PERF_TEST_"
+----
+
+With `RABBITMQ_STREAM_PERF_TEST_ENV_PREFIX="STREAM_PERF_TEST_"` defined, the tool looks for the `STREAM_PERF_TEST_BATCH_SIZE` environment variable, not `BATCH_SIZE`.
+
 ===== Logging
 
 The performance tool binary uses Logback with an internal configuration file.
diff --git a/src/main/java/com/rabbitmq/stream/perf/StreamPerfTest.java b/src/main/java/com/rabbitmq/stream/perf/StreamPerfTest.java
@@ -13,6 +13,9 @@
 // info@rabbitmq.com.
 package com.rabbitmq.stream.perf;
 
+import static com.rabbitmq.stream.perf.Utils.ENVIRONMENT_VARIABLE_LOOKUP;
+import static com.rabbitmq.stream.perf.Utils.ENVIRONMENT_VARIABLE_PREFIX;
+import static com.rabbitmq.stream.perf.Utils.OPTION_TO_ENVIRONMENT_VARIABLE;
 import static java.lang.String.format;
 
 import com.google.common.util.concurrent.RateLimiter;
@@ -74,6 +77,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -82,6 +86,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import picocli.CommandLine;
+import picocli.CommandLine.Model.CommandSpec;
 
 @CommandLine.Command(
     name = "stream-perf-test",
@@ -329,6 +334,12 @@ public class StreamPerfTest implements Callable<Integer> {
       defaultValue = "8080")
   private int monitoringPort;
 
+  @CommandLine.Option(
+      names = {"--environment-variables", "-env"},
+      description = "show usage with environment variables",
+      defaultValue = "false")
+  private boolean environmentVariables;
+
   private MetricsCollector metricsCollector;
   private PerformanceMetrics performanceMetrics;
   private List<Monitoring> monitorings;
@@ -422,10 +433,9 @@ private static boolean isTls(Collection<String> uris) {
 
   @Override
   public Integer call() throws Exception {
-    if (this.version) {
-      versionInformation(System.out);
-      System.exit(0);
-    }
+    maybeDisplayVersion();
+    maybeDisplayEnvironmentVariablesHelp();
+    overridePropertiesWithEnvironmentVariables();
 
     Codec codec = createCodec(this.codecClass);
 
@@ -808,6 +818,44 @@ public Integer call() throws Exception {
     return 0;
   }
 
+  private void overridePropertiesWithEnvironmentVariables() throws Exception {
+    Function<String, String> optionToEnvMappings =
+        OPTION_TO_ENVIRONMENT_VARIABLE
+            .andThen(ENVIRONMENT_VARIABLE_PREFIX)
+            .andThen(ENVIRONMENT_VARIABLE_LOOKUP);
+
+    Utils.assignValuesToCommand(this, optionToEnvMappings);
+    this.monitorings.forEach(
+        command -> {
+          try {
+            Utils.assignValuesToCommand(command, optionToEnvMappings);
+          } catch (Exception e) {
+            LOGGER.warn(
+                "Error while trying to assign environment variables to command {}",
+                command.getClass());
+          }
+        });
+  }
+
+  private void maybeDisplayEnvironmentVariablesHelp() {
+    if (this.environmentVariables) {
+      Collection<Object> commands = new ArrayList<>(this.monitorings.size() + 1);
+      commands.add(this);
+      commands.addAll(this.monitorings);
+      CommandSpec commandSpec = Utils.buildCommandSpec(commands.toArray());
+      CommandLine commandLine = new CommandLine(commandSpec);
+      CommandLine.usage(commandLine, System.out);
+      System.exit(0);
+    }
+  }
+
+  private void maybeDisplayVersion() {
+    if (this.version) {
+      versionInformation(System.out);
+      System.exit(0);
+    }
+  }
+
   private ShutdownService.CloseCallback closeStep(
       String message, ShutdownService.CloseCallback callback) {
     return new CloseCallback() {
diff --git a/src/main/java/com/rabbitmq/stream/perf/Utils.java b/src/main/java/com/rabbitmq/stream/perf/Utils.java
@@ -18,6 +18,8 @@
 import com.rabbitmq.stream.StreamCreator.LeaderLocator;
 import com.rabbitmq.stream.compression.Compression;
 import com.sun.management.OperatingSystemMXBean;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.security.cert.X509Certificate;
 import java.text.CharacterIterator;
@@ -27,8 +29,11 @@
 import java.time.format.DateTimeFormatter;
 import java.time.format.DateTimeParseException;
 import java.time.temporal.TemporalAccessor;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
@@ -38,6 +43,7 @@
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.function.LongSupplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -47,11 +53,38 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import picocli.CommandLine;
+import picocli.CommandLine.Command;
 import picocli.CommandLine.ITypeConverter;
+import picocli.CommandLine.Model.CommandSpec;
+import picocli.CommandLine.Model.OptionSpec;
+import picocli.CommandLine.Option;
 
 class Utils {
 
   static final X509TrustManager TRUST_EVERYTHING_TRUST_MANAGER = new TrustEverythingTrustManager();
+  static final Function<String, String> OPTION_TO_ENVIRONMENT_VARIABLE =
+      option -> {
+        if (option.startsWith("--")) {
+          return option.replace("--", "").replace('-', '_').toUpperCase(Locale.ENGLISH);
+        } else if (option.startsWith("-")) {
+          return option.substring(1).replace('-', '_').toUpperCase(Locale.ENGLISH);
+        } else {
+          return option.replace('-', '_').toUpperCase(Locale.ENGLISH);
+        }
+      };
+  static final Function<String, String> ENVIRONMENT_VARIABLE_PREFIX =
+      name -> {
+        String prefix = System.getenv("RABBITMQ_STREAM_PERF_TEST_ENV_PREFIX");
+        if (prefix == null || prefix.trim().isEmpty()) {
+          return name;
+        }
+        if (prefix.endsWith("_")) {
+          return prefix + name;
+        } else {
+          return prefix + "_" + name;
+        }
+      };
+  static final Function<String, String> ENVIRONMENT_VARIABLE_LOOKUP = name -> System.getenv(name);
   private static final LongSupplier TOTAL_MEMORY_SIZE_SUPPLIER;
   private static final Logger LOGGER = LoggerFactory.getLogger(Utils.class);
   private static final String RANGE_SEPARATOR_1 = "-";
@@ -166,6 +199,85 @@ private static void throwConversionException(String format, String... arguments)
     throw new CommandLine.TypeConversionException(String.format(format, (Object[]) arguments));
   }
 
+  static void assignValuesToCommand(Object command, Function<String, String> optionMapping)
+      throws Exception {
+    LOGGER.debug("Assigning values to command {}", command.getClass());
+    Collection<String> arguments = new ArrayList<>();
+    Collection<Field> fieldsToAssign = new ArrayList<>();
+    for (Field field : command.getClass().getDeclaredFields()) {
+      Option option = field.getAnnotation(Option.class);
+      if (option == null) {
+        LOGGER.debug("No option annotation for field {}", field.getName());
+        continue;
+      }
+      String longOption =
+          Arrays.stream(option.names())
+              .sorted(Comparator.comparingInt(String::length).reversed())
+              .findFirst()
+              .get();
+      LOGGER.debug("Looking up new value for option {}", longOption);
+      String newValue = optionMapping.apply(longOption);
+
+      LOGGER.debug(
+          "New value found for option {} (field {}): {}", longOption, field.getName(), newValue);
+      if (newValue == null) {
+        continue;
+      }
+      fieldsToAssign.add(field);
+      if (field.getType().equals(boolean.class) || field.getType().equals(Boolean.class)) {
+        if (Boolean.parseBoolean(newValue)) {
+          arguments.add(longOption);
+        }
+      } else {
+        arguments.add(longOption + " " + newValue);
+      }
+    }
+    if (fieldsToAssign.size() > 0) {
+      Constructor<?> defaultConstructor = command.getClass().getConstructor();
+      Object commandBuffer = defaultConstructor.newInstance();
+      String argumentsLine = String.join(" ", arguments);
+      LOGGER.debug("Arguments line with extra values: {}", argumentsLine);
+      String[] args = argumentsLine.split(" ");
+      commandBuffer = CommandLine.populateCommand(commandBuffer, args);
+      for (Field field : fieldsToAssign) {
+        field.setAccessible(true);
+        field.set(command, field.get(commandBuffer));
+      }
+    }
+  }
+
+  static CommandSpec buildCommandSpec(Object... commands) {
+    Object mainCommand = commands[0];
+    Command commandAnnotation = mainCommand.getClass().getAnnotation(Command.class);
+    CommandSpec spec = CommandSpec.create();
+    spec.name(commandAnnotation.name());
+    spec.mixinStandardHelpOptions(commandAnnotation.mixinStandardHelpOptions());
+    for (Object command : commands) {
+      for (Field f : command.getClass().getDeclaredFields()) {
+        Option annotation = f.getAnnotation(Option.class);
+        if (annotation == null) {
+          continue;
+        }
+        String name =
+            Arrays.stream(annotation.names())
+                .sorted(Comparator.comparingInt(String::length).reversed())
+                .findFirst()
+                .map(OPTION_TO_ENVIRONMENT_VARIABLE::apply)
+                .get();
+        spec.addOption(
+            OptionSpec.builder(name)
+                .type(f.getType())
+                .description(annotation.description())
+                .paramLabel("<" + name.replace("_", "-") + ">")
+                .defaultValue(annotation.defaultValue())
+                .showDefaultValue(annotation.showDefaultValue())
+                .build());
+      }
+    }
+
+    return spec;
+  }
+
   static class ByteCapacityTypeConverter implements CommandLine.ITypeConverter<ByteCapacity> {
 
     @Override
diff --git a/src/test/java/com/rabbitmq/stream/perf/UtilsTest.java b/src/test/java/com/rabbitmq/stream/perf/UtilsTest.java
@@ -28,8 +28,10 @@
 import com.rabbitmq.stream.perf.Utils.SniServerNamesConverter;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Random;
 import java.util.UUID;
 import java.util.function.BiFunction;
@@ -44,6 +46,9 @@
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import picocli.CommandLine;
+import picocli.CommandLine.Command;
+import picocli.CommandLine.Model.CommandSpec;
+import picocli.CommandLine.Option;
 import picocli.CommandLine.TypeConversionException;
 
 public class UtilsTest {
@@ -263,4 +268,58 @@ void streamsSpecifyOnlyOneStreamWhenStreamCountIsSpecified() {
     assertThatThrownBy(() -> Utils.streams("2", Arrays.asList("stream1", "stream2")))
         .isInstanceOf(IllegalArgumentException.class);
   }
+
+  @Test
+  void assignValuesToCommand() throws Exception {
+    TestCommand command = new TestCommand();
+    Map<String, String> mappings = new HashMap<>();
+    mappings.put("aaa", "42"); // takes only long options
+    mappings.put("b", "true");
+    mappings.put("offset", "first");
+    Utils.assignValuesToCommand(command, option -> mappings.get(option));
+    assertThat(command.a).isEqualTo(42);
+    assertThat(command.b).isTrue();
+    assertThat(command.c).isFalse();
+    assertThat(command.offsetSpecification).isEqualTo(OffsetSpecification.first());
+  }
+
+  @Test
+  void buildCommandSpec() {
+    CommandSpec spec = Utils.buildCommandSpec(new TestCommand());
+    assertThat(spec.optionsMap()).hasSize(4).containsKeys("AAA", "B", "C", "OFFSET");
+  }
+
+  @ParameterizedTest
+  @CsvSource({
+    "--uris,URIS",
+    "--stream-count,STREAM_COUNT",
+    "-sc,SC",
+    "--sub-entry-size, SUB_ENTRY_SIZE",
+    "-ses,SES"
+  })
+  void optionToEnvironmentVariable(String option, String envVariable) {
+    assertThat(Utils.OPTION_TO_ENVIRONMENT_VARIABLE.apply(option)).isEqualTo(envVariable);
+  }
+
+  @Command(name = "test-command")
+  static class TestCommand {
+    @Option(
+        names = {"aaa", "a"},
+        defaultValue = "10")
+    private int a = 10;
+
+    @Option(names = "b", defaultValue = "false")
+    private boolean b = false;
+
+    @Option(names = "c", defaultValue = "false")
+    private boolean c = false;
+
+    @CommandLine.Option(
+        names = {"offset"},
+        defaultValue = "next",
+        converter = Utils.OffsetSpecificationTypeConverter.class)
+    private OffsetSpecification offsetSpecification = OffsetSpecification.next();
+
+    public TestCommand() {}
+  }
 }
