[ISSUE apache#3487] Benchmark supports custom AK/SK

Signed-off-by: zhangyang21 <zhangyang21@xiaomi.com>
GenerousMan · Nov 17, 2021 · dc252cf · dc252cf
1 parent cead705
commit dc252cf
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 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
@@ -23,11 +23,15 @@
 
 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
@@ -123,7 +123,12 @@ public void run() {
             }
         }, 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 static Options buildCommandlineOptions(final Options options) {
         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
@@ -126,7 +126,12 @@ public void run() {
             }
         }, 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 static Options buildCommandlineOptions(final Options options) {
         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
@@ -31,6 +31,7 @@
 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 void run() {
             }
         }, 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 static Options buildCommandlineOptions(final Options options) {
         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);