From b7c405df9a193694368abdf0d2ba8d4ca37647c2 Mon Sep 17 00:00:00 2001
From: Shawn <506895667@qq.com>
Date: Mon, 9 Sep 2024 19:03:30 +0800
Subject: [PATCH] feat: support %i vars as username opt

---
 src/emqtt_bench.erl | 88 ++++++++++++++++++++++++++-------------------
 1 file changed, 51 insertions(+), 37 deletions(-)

diff --git a/src/emqtt_bench.erl b/src/emqtt_bench.erl
index 8435b9e..7fd8eb4 100644
--- a/src/emqtt_bench.erl
+++ b/src/emqtt_bench.erl
@@ -252,7 +252,7 @@
          {interval, $i, "interval", {integer, 10},
           "interval of connecting to the broker"},
          {username, $u, "username", string,
-          "username for connecting to server"},
+          "username for connecting to server, support %i, %(i/2), %(rand(100)) variables"},
          {password, $P, "password", string,
           "password for connecting to server"},
          {keepalive, $k, "keepalive", {integer, 300},
@@ -299,6 +299,25 @@
          }
         ]).
 
+-define(COUNTER_NAMES,
+    [ publish_latency
+    , recv
+    , sub
+    , sub_fail
+    , pub
+    , pub_fail
+    , pub_overrun
+    , pub_succ
+    , connect_succ
+    , connect_fail
+    , connect_retried
+    , reconnect_succ
+    , unreachable
+    , connection_refused
+    , connection_timeout
+    , connection_idle
+    ]).
+
 -define(cnt_map, cnt_map).
 -define(hdr_cnt64, "cnt64").
 -define(hdr_ts, "ts").
@@ -634,6 +653,7 @@ run(Parent, I, N, PubSub, Opts0, AddrList, HostList) ->
     run(Parent, I + 1, N, PubSub, Opts, AddrList, HostList).
 
 connect(Parent, N, PubSub, Opts) ->
+    _ = erlang:put(conn_id, N),
     process_flag(trap_exit, true),
     rand:seed(exsplus, erlang:timestamp()),
     Prometheus = lists:member(prometheus, Opts),
@@ -952,8 +972,9 @@ mqtt_opts([{version, 4}|Opts], Acc) ->
     mqtt_opts(Opts, [{proto_ver, v4}|Acc]);
 mqtt_opts([{version, 5}|Opts], Acc) ->
     mqtt_opts(Opts, [{proto_ver, v5}|Acc]);
-mqtt_opts([{username, Username}|Opts], Acc) ->
-    mqtt_opts(Opts, [{username, list_to_binary(Username)}|Acc]);
+mqtt_opts([{username, Username0}|Opts], Acc) ->
+    Username = parse_expr(Username0),
+    mqtt_opts(Opts, [{username, bin(Username)}|Acc]);
 mqtt_opts([{password, Password}|Opts], Acc) ->
     mqtt_opts(Opts, [{password, list_to_binary(Password)}|Acc]);
 mqtt_opts([{keepalive, I}|Opts], Acc) ->
@@ -981,6 +1002,31 @@ mqtt_opts([{reconnect, Reconnect}|Opts], Acc) ->
 mqtt_opts([_|Opts], Acc) ->
     mqtt_opts(Opts, Acc).
 
+eval_expr("random(" ++ Str) ->
+    case lists:last(Str) of
+        $) ->
+            Int = strip_last_char(Str),
+            rand:uniform(list_to_integer(Int));
+        _ -> Str
+    end.
+
+parse_expr("%i") ->
+    erlang:get(conn_id);
+parse_expr("%(i div 2)") ->
+    erlang:get(conn_id) div 2;
+parse_expr("%(" ++ Str) ->
+    case lists:last(Str) of
+        $) ->
+            eval_expr(strip_last_char(Str));
+        _ ->
+            throw({invalid_expression, "%(" ++ Str})
+    end;
+parse_expr(Str) ->
+    Str.
+
+strip_last_char(Str) ->
+    lists:sublist(Str, 1, length(Str) - 1).
+
 tcp_opts(Opts) ->
     tcp_opts(Opts, []).
 tcp_opts([], Acc) ->
@@ -1250,23 +1296,7 @@ prepare_for_quic(Opts)->
 
 -spec counters() -> {atom(), integer()}.
 counters() ->
-   Names = [ publish_latency
-           , recv
-           , sub
-           , sub_fail
-           , pub
-           , pub_fail
-           , pub_overrun
-           , pub_succ
-           , connect_succ
-           , connect_fail
-           , connect_retried
-           , reconnect_succ
-           , unreachable
-           , connection_refused
-           , connection_timeout
-           , connection_idle
-           ],
+   Names = ?COUNTER_NAMES,
    Idxs = lists:seq(2, length(Names) + 1),
    lists:zip(Names, Idxs).
 
@@ -1281,23 +1311,7 @@ maybe_init_prometheus(true) ->
                    prometheus_summary],
     application:set_env(prometheus, collectors, Collectors),
     {ok, _} = application:ensure_all_started(prometheus),
-    Counters = [ publish_latency
-               , recv
-               , sub
-               , sub_fail
-               , pub
-               , pub_fail
-               , pub_overrun
-               , pub_succ
-               , connect_succ
-               , connect_fail
-               , connect_retried
-               , reconnect_succ
-               , unreachable
-               , connection_refused
-               , connection_timeout
-               , connection_idle
-               ],
+    Counters = ?COUNTER_NAMES,
     lists:foreach(
       fun(Cnt) ->
               prometheus_counter:declare([{name, Cnt}, {help, atom_to_list(Cnt)}])