Skip to content

Commit 988a2bf

Browse files
acogoluegnesacogoluegnes
committed
Support environment variables as arguments
This commit allows to use environment variables to configure PerfTest. Environment variables are prioritary over command line arguments. The name of the environment variables is a snake case version of the long command line argument name, e.g. --skip-binding-queues becomes SKIP_BINDING_QUEUES. A prefix can be used by setting the RABBITMQ_PERF_TEST_ENV_PREFIX environment variable. [#157994032] Fixes #97
1 parent b7c536d commit 988a2bf

File tree

4 files changed

+317
-25
lines changed

4 files changed

+317
-25
lines changed

‎src/main/java/com/rabbitmq/perf/CommandLineProxy.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
// Copyright (c) 2018-Present Pivotal Software, Inc. All rights reserved.
2+
//
3+
// This software, the RabbitMQ Java client library, is triple-licensed under the
4+
// Mozilla Public License 1.1 ("MPL"), the GNU General Public License version 2
5+
// ("GPL") and the Apache License version 2 ("ASL"). For the MPL, please see
6+
// LICENSE-MPL-RabbitMQ. For the GPL, please see LICENSE-GPL2. For the ASL,
7+
// please see LICENSE-APACHE2.
8+
//
9+
// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
10+
// either express or implied. See the LICENSE file for specific language governing
11+
// rights and limitations of this software.
12+
//
13+
// If you have any questions regarding licensing, please contact us at
14+
// [email protected].
15+
16+
package com.rabbitmq.perf;
17+
18+
import org.apache.commons.cli.CommandLine;
19+
import org.apache.commons.cli.Option;
20+
import org.apache.commons.cli.Options;
21+
22+
import java.util.function.Function;
23+
import java.util.function.Supplier;
24+
25+
import static java.lang.String.valueOf;
26+
27+
/**
28+
* A proxy to add behavior around a {@link CommandLine}.
29+
* It implements only the few methods used in {@link PerfTest}.
30+
* This class has been introduced to easily make environment
31+
* variables override or directly set arguments.
32+
*
33+
* {@link CommandLine} doesn't implement any interface nor
34+
* can be subclassed, that's why this proxy trick is used.
35+
*/
36+
public class CommandLineProxy {
37+
38+
private final CommandLine delegate;
39+
40+
private final Function<String, String> argumentLookup;
41+
42+
public CommandLineProxy(final Options options, CommandLine delegate, Function<String, String> argumentLookup) {
43+
this.delegate = delegate;
44+
Function<String, String> optionToLongOption = option -> {
45+
for (Object optObj : options.getOptions()) {
46+
Option opt = (Option) optObj;
47+
if (opt.getOpt().equals(option)) {
48+
return opt.getLongOpt();
49+
}
50+
}
51+
return null;
52+
};
53+
this.argumentLookup = optionToLongOption
54+
.andThen(longOption -> longOption == null ? null : argumentLookup.apply(longOption));
55+
}
56+
57+
public boolean hasOption(char opt) {
58+
return override(valueOf(opt), () -> delegate.hasOption(opt), Boolean.class);
59+
}
60+
61+
public boolean hasOption(String opt) {
62+
return override(opt, () -> delegate.hasOption(opt), Boolean.class);
63+
}
64+
65+
public String getOptionValue(char opt, String def) {
66+
return override(valueOf(opt), () -> delegate.getOptionValue(opt, def), String.class);
67+
}
68+
69+
public String getOptionValue(String opt, String def) {
70+
return override(opt, () -> delegate.getOptionValue(opt, def), String.class);
71+
}
72+
73+
public String[] getOptionValues(char opt) {
74+
return override(valueOf(opt), () -> delegate.getOptionValues(opt), String[].class);
75+
}
76+
77+
public String getOptionValue(char opt) {
78+
return override(valueOf(opt), () -> delegate.getOptionValue(opt), String.class);
79+
}
80+
81+
@SuppressWarnings("unchecked")
82+
private <T> T override(String opt, Supplier<T> argumentValue, Class<T> type) {
83+
String value = argumentLookup.apply(opt);
84+
if (value == null) {
85+
return argumentValue.get();
86+
} else {
87+
if (Boolean.class.equals(type)) {
88+
return (T) Boolean.valueOf(value);
89+
} else if (String[].class.equals(type)) {
90+
return (T) value.split(",");
91+
} else {
92+
return (T) value;
93+
}
94+
}
95+
}
96+
}

‎src/main/java/com/rabbitmq/perf/PerfTest.java‎

Lines changed: 56 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,15 @@
2020
import java.security.NoSuchAlgorithmException;
2121
import java.text.SimpleDateFormat;
2222
import java.util.Calendar;
23-
import java.util.Collections;
2423
import java.util.HashMap;
2524
import java.util.List;
2625
import java.util.Locale;
2726
import java.util.Map;
2827
import java.util.Properties;
29-
import java.util.concurrent.Executors;
3028
import java.util.concurrent.SynchronousQueue;
3129
import java.util.concurrent.ThreadPoolExecutor;
3230
import java.util.concurrent.TimeUnit;
31+
import java.util.function.Function;
3332

3433
import com.rabbitmq.client.impl.ClientVersion;
3534
import com.rabbitmq.client.impl.nio.NioParams;
@@ -50,6 +49,7 @@
5049

5150
import static java.lang.String.format;
5251
import static java.util.Arrays.asList;
52+
import static java.util.Collections.singletonList;
5353

5454
public class PerfTest {
5555

@@ -60,7 +60,8 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
6060
Options options = getOptions();
6161
CommandLineParser parser = new GnuParser();
6262
try {
63-
CommandLine cmd = parser.parse(options, args);
63+
CommandLine rawCmd = parser.parse(options, args);
64+
CommandLineProxy cmd = new CommandLineProxy(options, rawCmd, perfTestOptions.argumentLookup);
6465

6566
if (cmd.hasOption('?')) {
6667
usage(options);
@@ -79,8 +80,8 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
7980
String exchangeName = getExchangeName(cmd, exchangeType);
8081
String queueNames = strArg(cmd, 'u', null);
8182
String routingKey = strArg(cmd, 'k', null);
82-
boolean randomRoutingKey = cmd.hasOption('K');
83-
boolean skipBindingQueues= cmd.hasOption("sb");
83+
boolean randomRoutingKey = hasOption(cmd, "K");
84+
boolean skipBindingQueues= hasOption(cmd,"sb");
8485
int samplingInterval = intArg(cmd, 'i', 1);
8586
float producerRateLimit = floatArg(cmd, 'r', 0.0f);
8687
float consumerRateLimit = floatArg(cmd, 'R', 0.0f);
@@ -92,12 +93,12 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
9293
int consumerTxSize = intArg(cmd, 'n', 0);
9394
long confirm = intArg(cmd, 'c', -1);
9495
int confirmTimeout = intArg(cmd, "ct", 30);
95-
boolean autoAck = cmd.hasOption('a');
96+
boolean autoAck = hasOption(cmd,"a");
9697
int multiAckEvery = intArg(cmd, 'A', 0);
9798
int channelPrefetch = intArg(cmd, 'Q', 0);
9899
int consumerPrefetch = intArg(cmd, 'q', 0);
99100
int minMsgSize = intArg(cmd, 's', 0);
100-
boolean slowStart = cmd.hasOption('S');
101+
boolean slowStart = hasOption(cmd, "S");
101102
int timeLimit = intArg(cmd, 'z', 0);
102103
int producerMsgCount = intArg(cmd, 'C', 0);
103104
int consumerMsgCount = intArg(cmd, 'D', 0);
@@ -106,17 +107,17 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
106107
int heartbeat = intArg(cmd, 'b', 0);
107108
String bodyFiles = strArg(cmd, 'B', null);
108109
String bodyContentType = strArg(cmd, 'T', null);
109-
boolean predeclared = cmd.hasOption('p');
110-
boolean legacyMetrics = cmd.hasOption('l');
110+
boolean predeclared = hasOption(cmd, "p");
111+
boolean legacyMetrics = hasOption(cmd, "l");
111112
boolean autoDelete = boolArg(cmd, "ad", true);
112-
boolean useMillis = cmd.hasOption("ms");
113-
boolean saslExternal = cmd.hasOption("se");
113+
boolean useMillis = hasOption(cmd,"ms");
114+
boolean saslExternal = hasOption(cmd, "se");
114115
String queueArgs = strArg(cmd, "qa", null);
115116
int consumerLatencyInMicroseconds = intArg(cmd, 'L', 0);
116117
int heartbeatSenderThreads = intArg(cmd, "hst", -1);
117118
String messageProperties = strArg(cmd, "mp", null);
118119
int routingKeyCacheSize = intArg(cmd, "rkcs", 0);
119-
boolean exclusive = cmd.hasOption('E');
120+
boolean exclusive = hasOption(cmd, "E");
120121
int publishingIntervalInSeconds = intArg(cmd, "P", -1);
121122
int producerRandomStartDelayInSeconds = intArg(cmd, "prsd", -1);
122123
int producerSchedulingThreads = intArg(cmd, "pst", -1);
@@ -145,7 +146,7 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
145146
}
146147
uris = asList(urisArray);
147148
} else {
148-
uris = Collections.singletonList(uri);
149+
uris = singletonList(uri);
149150
}
150151

151152
//setup
@@ -254,7 +255,7 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
254255
}
255256
}
256257

257-
private static ConnectionFactory configureNioIfRequested(CommandLine cmd, ConnectionFactory factory) {
258+
private static ConnectionFactory configureNioIfRequested(CommandLineProxy cmd, ConnectionFactory factory) {
258259
int nbThreads = intArg(cmd, "niot", -1);
259260
int executorSize = intArg(cmd, "niotp", -1);
260261
if (nbThreads > 0 || executorSize > 0) {
@@ -313,9 +314,9 @@ public static void main(String[] args) {
313314
main(args, new PerfTestOptions().setSystemExiter(new JvmSystemExiter()).setSkipSslContextConfiguration(false));
314315
}
315316

316-
private static SSLContext getSslContextIfNecessary(CommandLine cmd, Properties systemProperties) throws NoSuchAlgorithmException {
317+
private static SSLContext getSslContextIfNecessary(CommandLineProxy cmd, Properties systemProperties) throws NoSuchAlgorithmException {
317318
SSLContext sslContext = null;
318-
if (cmd.hasOption("useDefaultSslContext")) {
319+
if (hasOption(cmd, "udsc") || hasOption(cmd,"useDefaultSslContext")) {
319320
LOGGER.info("Using default SSL context as per command line option");
320321
sslContext = SSLContext.getDefault();
321322
}
@@ -338,7 +339,7 @@ private static void usage(Options options) {
338339
formatter.printHelp("<program>", options);
339340
}
340341

341-
private static Options getOptions() {
342+
static Options getOptions() {
342343
Options options = new Options();
343344
options.addOption(new Option("?", "help", false,"show usage"));
344345
options.addOption(new Option("d", "id", true, "test ID"));
@@ -417,38 +418,42 @@ private static Options getOptions() {
417418
return options;
418419
}
419420

420-
private static String strArg(CommandLine cmd, char opt, String def) {
421+
private static String strArg(CommandLineProxy cmd, char opt, String def) {
421422
return cmd.getOptionValue(opt, def);
422423
}
423424

424-
private static String strArg(CommandLine cmd, String opt, String def) {
425+
private static String strArg(CommandLineProxy cmd, String opt, String def) {
425426
return cmd.getOptionValue(opt, def);
426427
}
427428

428-
private static int intArg(CommandLine cmd, char opt, int def) {
429+
private static int intArg(CommandLineProxy cmd, char opt, int def) {
429430
return Integer.parseInt(cmd.getOptionValue(opt, Integer.toString(def)));
430431
}
431432

432-
private static int intArg(CommandLine cmd, String opt, int def) {
433+
private static int intArg(CommandLineProxy cmd, String opt, int def) {
433434
return Integer.parseInt(cmd.getOptionValue(opt, Integer.toString(def)));
434435
}
435436

436-
private static float floatArg(CommandLine cmd, char opt, float def) {
437+
private static float floatArg(CommandLineProxy cmd, char opt, float def) {
437438
return Float.parseFloat(cmd.getOptionValue(opt, Float.toString(def)));
438439
}
439440

440-
private static boolean boolArg(CommandLine cmd, String opt, boolean def) {
441+
private static boolean boolArg(CommandLineProxy cmd, String opt, boolean def) {
441442
return Boolean.parseBoolean(cmd.getOptionValue(opt, Boolean.toString(def)));
442443
}
443444

444-
private static List<?> lstArg(CommandLine cmd, char opt) {
445+
private static List<?> lstArg(CommandLineProxy cmd, char opt) {
445446
String[] vals = cmd.getOptionValues(opt);
446447
if (vals == null) {
447448
vals = new String[] {};
448449
}
449450
return asList(vals);
450451
}
451452

453+
private static boolean hasOption(CommandLineProxy cmd, String opt) {
454+
return cmd.hasOption(opt);
455+
}
456+
452457
private static Map<String, Object> convertKeyValuePairs(String arg) {
453458
if (arg == null || arg.trim().isEmpty()) {
454459
return null;
@@ -465,7 +470,7 @@ private static Map<String, Object> convertKeyValuePairs(String arg) {
465470
return properties;
466471
}
467472

468-
private static String getExchangeName(CommandLine cmd, String def) {
473+
private static String getExchangeName(CommandLineProxy cmd, String def) {
469474
String exchangeName = null;
470475
if (cmd.hasOption('e')) {
471476
exchangeName = cmd.getOptionValue('e');
@@ -648,6 +653,10 @@ public static class PerfTestOptions {
648653

649654
private boolean skipSslContextConfiguration = false;
650655

656+
private Function<String, String> argumentLookup = LONG_OPTION_TO_ENVIRONMENT_VARIABLE
657+
.andThen(ENVIRONMENT_VARIABLE_PREFIX)
658+
.andThen(ENVIRONMENT_VARIABLE_LOOKUP);
659+
651660
public PerfTestOptions setSystemExiter(SystemExiter systemExiter) {
652661
this.systemExiter = systemExiter;
653662
return this;
@@ -657,6 +666,11 @@ public PerfTestOptions setSkipSslContextConfiguration(boolean skipSslContextConf
657666
this.skipSslContextConfiguration = skipSslContextConfiguration;
658667
return this;
659668
}
669+
670+
public PerfTestOptions setArgumentLookup(Function<String, String> argumentLookup) {
671+
this.argumentLookup = argumentLookup;
672+
return this;
673+
}
660674
}
661675

662676
/**
@@ -682,4 +696,21 @@ public void exit(int status) {
682696

683697
}
684698

699+
static Function<String, String> LONG_OPTION_TO_ENVIRONMENT_VARIABLE = option ->
700+
option.replace('-', '_').toUpperCase(Locale.ENGLISH);
701+
702+
static Function<String, String> ENVIRONMENT_VARIABLE_PREFIX = name -> {
703+
String prefix = System.getenv("RABBITMQ_PERF_TEST_ENV_PREFIX");
704+
if (prefix == null || prefix.trim().isEmpty()) {
705+
return name;
706+
}
707+
if (prefix.endsWith("_")) {
708+
return prefix + name;
709+
} else {
710+
return prefix + "_" + name;
711+
}
712+
};
713+
714+
static Function<String, String> ENVIRONMENT_VARIABLE_LOOKUP = name -> System.getenv(name);
715+
685716
}

0 commit comments

Comments
 (0)