From 98fdf4eea0390ac9410b01715f43293e491ea506 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 | 70 ++++++++++++++++++++++++++++++++-------------
 1 file changed, 50 insertions(+), 20 deletions(-)

diff --git a/src/emqtt_bench.erl b/src/emqtt_bench.erl
index 7607cdd..852ec71 100644
--- a/src/emqtt_bench.erl
+++ b/src/emqtt_bench.erl
@@ -236,7 +236,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},
@@ -275,6 +275,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").
@@ -590,6 +609,7 @@ run(Parent, N, PubSub, Opts0, AddrList, HostList) ->
     run(Parent, N-1, PubSub, Opts, AddrList, HostList).
 
 connect(Parent, N, PubSub, Opts) ->
+    _ = erlang:put(conn_id, N),
     process_flag(trap_exit, true),
     rand:seed(exsplus, erlang:timestamp()),
     MRef = case proplists:get_value(publish_signal_pid, Opts) of
@@ -892,8 +912,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) ->
@@ -921,6 +942,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) ->
@@ -1188,23 +1234,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).