From 4bb99e656b1acc99d3a953f84ae0abcb74737af5 Mon Sep 17 00:00:00 2001 From: zhangyang Date: Wed, 17 Nov 2021 13:56:46 +0800 Subject: [PATCH] [ISSUE #3487] Benchmark supports custom AK/SK Signed-off-by: zhangyang21 --- .../rocketmq/example/benchmark/AclClient.java | 12 ++++++++---- .../rocketmq/example/benchmark/Consumer.java | 15 ++++++++++++++- .../rocketmq/example/benchmark/Producer.java | 15 ++++++++++++++- .../example/benchmark/TransactionProducer.java | 17 ++++++++++++++++- 4 files changed, 52 insertions(+), 7 deletions(-) diff --git a/example/src/main/java/org/apache/rocketmq/example/benchmark/AclClient.java b/example/src/main/java/org/apache/rocketmq/example/benchmark/AclClient.java index 04ef5d5ef6..b3d6fb47e9 100644 --- a/example/src/main/java/org/apache/rocketmq/example/benchmark/AclClient.java +++ b/example/src/main/java/org/apache/rocketmq/example/benchmark/AclClient.java @@ -23,11 +23,15 @@ import org.apache.rocketmq.remoting.RPCHook; public class AclClient { - private static final String ACL_ACCESS_KEY = "rocketmq2"; + public static final String ACL_ACCESS_KEY = "rocketmq2"; - private static final String ACL_SECRET_KEY = "12345678"; + public static final String ACL_SECRET_KEY = "12345678"; - static RPCHook getAclRPCHook() { - return new AclClientRPCHook(new SessionCredentials(ACL_ACCESS_KEY,ACL_SECRET_KEY)); + public static RPCHook getAclRPCHook() { + return getAclRPCHook(ACL_ACCESS_KEY, ACL_SECRET_KEY); + } + + public static RPCHook getAclRPCHook(String ak, String sk) { + return new AclClientRPCHook(new SessionCredentials(ak, sk)); } } diff --git a/example/src/main/java/org/apache/rocketmq/example/benchmark/Consumer.java b/example/src/main/java/org/apache/rocketmq/example/benchmark/Consumer.java index 7d26509d2b..d08795d6b0 100644 --- a/example/src/main/java/org/apache/rocketmq/example/benchmark/Consumer.java +++ b/example/src/main/java/org/apache/rocketmq/example/benchmark/Consumer.java @@ -123,7 +123,12 @@ public class Consumer { } }, 10000, 10000, TimeUnit.MILLISECONDS); - RPCHook rpcHook = aclEnable ? AclClient.getAclRPCHook() : null; + RPCHook rpcHook = null; + if (aclEnable) { + String ak = commandLine.hasOption("ak") ? String.valueOf(commandLine.getOptionValue("ak")) : AclClient.ACL_ACCESS_KEY; + String sk = commandLine.hasOption("sk") ? String.valueOf(commandLine.getOptionValue("sk")) : AclClient.ACL_SECRET_KEY; + rpcHook = AclClient.getAclRPCHook(ak, sk); + } DefaultMQPushConsumer consumer = new DefaultMQPushConsumer(group, rpcHook, new AllocateMessageQueueAveragely(), msgTraceEnable, null); if (commandLine.hasOption('n')) { String ns = commandLine.getOptionValue('n'); @@ -218,6 +223,14 @@ public class Consumer { opt.setRequired(false); options.addOption(opt); + opt = new Option("ak", "accessKey", true, "Acl access key, Default: 12345678"); + opt.setRequired(false); + options.addOption(opt); + + opt = new Option("sk", "secretKey", true, "Acl secret key, Default: rocketmq2"); + opt.setRequired(false); + options.addOption(opt); + return options; } diff --git a/example/src/main/java/org/apache/rocketmq/example/benchmark/Producer.java b/example/src/main/java/org/apache/rocketmq/example/benchmark/Producer.java index feb25cd910..c32e00e97b 100644 --- a/example/src/main/java/org/apache/rocketmq/example/benchmark/Producer.java +++ b/example/src/main/java/org/apache/rocketmq/example/benchmark/Producer.java @@ -126,7 +126,12 @@ public class Producer { } }, 10000, 10000, TimeUnit.MILLISECONDS); - RPCHook rpcHook = aclEnable ? AclClient.getAclRPCHook() : null; + RPCHook rpcHook = null; + if (aclEnable) { + String ak = commandLine.hasOption("ak") ? String.valueOf(commandLine.getOptionValue("ak")) : AclClient.ACL_ACCESS_KEY; + String sk = commandLine.hasOption("sk") ? String.valueOf(commandLine.getOptionValue("sk")) : AclClient.ACL_SECRET_KEY; + rpcHook = AclClient.getAclRPCHook(ak, sk); + } final DefaultMQProducer producer = new DefaultMQProducer("benchmark_producer", rpcHook, msgTraceEnable, null); producer.setInstanceName(Long.toString(System.currentTimeMillis())); @@ -277,6 +282,14 @@ public class Producer { opt.setRequired(false); options.addOption(opt); + opt = new Option("ak", "accessKey", true, "Acl access key, Default: 12345678"); + opt.setRequired(false); + options.addOption(opt); + + opt = new Option("sk", "secretKey", true, "Acl secret key, Default: rocketmq2"); + opt.setRequired(false); + options.addOption(opt); + opt = new Option("q", "messageQuantity", true, "Send message quantity, Default: 0, running forever"); opt.setRequired(false); options.addOption(opt); diff --git a/example/src/main/java/org/apache/rocketmq/example/benchmark/TransactionProducer.java b/example/src/main/java/org/apache/rocketmq/example/benchmark/TransactionProducer.java index 767a96b40b..5e2f287e44 100644 --- a/example/src/main/java/org/apache/rocketmq/example/benchmark/TransactionProducer.java +++ b/example/src/main/java/org/apache/rocketmq/example/benchmark/TransactionProducer.java @@ -31,6 +31,7 @@ import org.apache.rocketmq.client.producer.TransactionMQProducer; import org.apache.rocketmq.common.message.Message; import org.apache.rocketmq.common.message.MessageConst; import org.apache.rocketmq.common.message.MessageExt; +import org.apache.rocketmq.remoting.RPCHook; import org.apache.rocketmq.srvutil.ServerUtil; import java.io.UnsupportedEncodingException; @@ -128,11 +129,17 @@ public class TransactionProducer { } }, 10000, 10000, TimeUnit.MILLISECONDS); + RPCHook rpcHook = null; + if (config.aclEnable) { + String ak = commandLine.hasOption("ak") ? String.valueOf(commandLine.getOptionValue("ak")) : AclClient.ACL_ACCESS_KEY; + String sk = commandLine.hasOption("sk") ? String.valueOf(commandLine.getOptionValue("sk")) : AclClient.ACL_SECRET_KEY; + rpcHook = AclClient.getAclRPCHook(ak, sk); + } final TransactionListener transactionCheckListener = new TransactionListenerImpl(statsBenchmark, config); final TransactionMQProducer producer = new TransactionMQProducer( null, "benchmark_transaction_producer", - config.aclEnable ? AclClient.getAclRPCHook() : null, + rpcHook, config.msgTraceEnable, null); producer.setInstanceName(Long.toString(System.currentTimeMillis())); @@ -268,6 +275,14 @@ public class TransactionProducer { opt.setRequired(false); options.addOption(opt); + opt = new Option("ak", "accessKey", true, "Acl access key, Default: 12345678"); + opt.setRequired(false); + options.addOption(opt); + + opt = new Option("sk", "secretKey", true, "Acl secret key, Default: rocketmq2"); + opt.setRequired(false); + options.addOption(opt); + opt = new Option("m", "msgTraceEnable", true, "Message Trace Enable, Default: false"); opt.setRequired(false); options.addOption(opt);