Skip to content

Commit f2f1afb

Browse files
lukebakkenlukebakken
authored
Merge pull request #98 from rabbitmq/rabbitmq-perf-test-env-variables-arguments
Support environment variables as arguments
2 parents b7c536d + 5d0da8e commit f2f1afb

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+
public 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+
public 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)