yuanbiao
/
emqx
1
0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

test(emqx_limiter): improve the code coverage of RateLimit to 87%

This commit is contained in:
lafirest 2022-02-17 14:48:04 +08:00
parent 6f1431d283
commit 3c92906405
4 changed files with 218 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
apps/emqx/src/emqx_limiter/src/emqx_htb_limiter.erl
View File

@ -334,7 +334,7 @@ make_retry_context(Fun, Diff) ->
-spec try_restore(retry_context(Limiter), Limiter) -> Limiter
when Limiter :: limiter().
try_restore(#{need := Need, diff := Diff},
#{tokens := Tokens, capcacity := Capacity, bucket := Bucket} = Limiter) ->
#{tokens := Tokens, capacity := Capacity, bucket := Bucket} = Limiter) ->
Back = Need - Diff,
Tokens2 = erlang:min(Capacity, Back + Tokens),
emqx_limiter_bucket_ref:try_restore(Back, Bucket),

2
apps/emqx/src/emqx_limiter/src/emqx_limiter_bucket_ref.erl
View File

@ -83,7 +83,7 @@ try_restore(_, infinity) ->
ok;
try_restore(Inc, #{counter := Counter, index := Index}) ->
case counters:get(Counter, Index) of
Tokens when Tokens < 0 ->
Tokens when Tokens =< 0 ->
counters:add(Counter, Index, Inc);
_ ->
ok

23
apps/emqx/src/emqx_limiter/src/emqx_limiter_server.erl
View File

@ -33,7 +33,7 @@
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
terminate/2, code_change/3, format_status/2]).
-export([ start_link/1, connect/2, info/2
-export([ start_link/1, connect/2, info/1
, name/1, get_initial_val/1]).
-type root() :: #{ rate := rate() %% number of tokens generated per period
@ -83,7 +83,7 @@
-type decimal() :: emqx_limiter_decimal:decimal().
-type index() :: pos_integer().
-define(CALL(Type, Msg), gen_server:call(name(Type), {?FUNCTION_NAME, Msg})).
-define(CALL(Type), gen_server:call(name(Type), ?FUNCTION_NAME)).
-define(OVERLOAD_MIN_ALLOC, 0.3). %% minimum coefficient for overloaded limiter
-export_type([index/0]).
@ -123,9 +123,9 @@ connect(Type, BucketName) when is_atom(BucketName) ->
connect(Type, Names) ->
connect(Type, maps:get(Type, Names, default)).
-spec info(limiter_type(), atom()) -> term().
info(Type, Info) ->
?CALL(Type, Info).
-spec info(limiter_type()) -> state().
info(Type) ->
?CALL(Type).
-spec name(limiter_type()) -> atom().
name(Type) ->
@ -183,6 +183,9 @@ init([Type]) ->
{noreply, NewState :: term(), hibernate} |
{stop, Reason :: term(), Reply :: term(), NewState :: term()} |
{stop, Reason :: term(), NewState :: term()}.
handle_call(info, _From, State) ->
{reply, State, State};
handle_call(Req, _From, State) ->
?SLOG(error, #{msg => "unexpected_call", call => Req}),
{reply, ignored, State}.
@ -375,11 +378,11 @@ maybe_burst(#{buckets := Buckets,
index := Index,
zone := Zone} = maps:get(Id, Nodes),
case counters:get(Counter, Index) of
Any when Any =< 0 ->
Group = maps:get(Zone, Groups, []),
maps:put(Zone, [Id | Group], Groups);
_ ->
Groups
Any when Any =< 0 ->
Group = maps:get(Zone, Groups, []),
maps:put(Zone, [Id | Group], Groups);
_ ->
Groups
end
end,

221
apps/emqx/test/emqx_ratelimiter_SUITE.erl
View File

@ -112,12 +112,63 @@ base_conf() ->
%%--------------------------------------------------------------------
%% Test Cases Bucket Level
%%--------------------------------------------------------------------
t_consume(_) ->
Cfg = fun(Cfg) ->
Cfg#{rate := 100,
capacity := 100,
initial := 100,
max_retry_time := 1000,
failure_strategy := force}
end,
Case = fun() ->
Client = connect(default),
{ok, L2} = emqx_htb_limiter:consume(50, Client),
{ok, _L3} = emqx_htb_limiter:consume(150, L2)
end,
with_per_client(default, Cfg, Case).
t_retry(_) ->
Cfg = fun(Cfg) ->
Cfg#{rate := 50,
capacity := 200,
initial := 0,
max_retry_time := 1000,
failure_strategy := force}
end,
Case = fun() ->
Client = connect(default),
{ok, Client} = emqx_htb_limiter:retry(Client),
{_, _, Retry, L2} = emqx_htb_limiter:check(150, Client),
L3 = emqx_htb_limiter:set_retry(Retry, L2),
timer:sleep(500),
{ok, _L4} = emqx_htb_limiter:retry(L3)
end,
with_per_client(default, Cfg, Case).
t_restore(_) ->
Cfg = fun(Cfg) ->
Cfg#{rate := 1,
capacity := 200,
initial := 50,
max_retry_time := 100,
failure_strategy := force}
end,
Case = fun() ->
Client = connect(default),
{_, _, Retry, L2} = emqx_htb_limiter:check(150, Client),
timer:sleep(200),
{ok, L3} = emqx_htb_limiter:check(Retry, L2),
Avaiable = emqx_htb_limiter:available(L3),
?assert(Avaiable >= 50)
end,
with_per_client(default, Cfg, Case).
t_max_retry_time(_) ->
Cfg = fun(Cfg) ->
Cfg#{rate := 1,
capacity := 1,
max_retry_time := 500,
failure_strategy := drop}
Cfg#{rate := 1,
capacity := 1,
max_retry_time := 500,
failure_strategy := drop}
end,
Case = fun() ->
Client = connect(default),
@ -186,6 +237,26 @@ t_infinity_client(_) ->
end,
with_bucket(default, Fun, Case).
t_try_restore_agg(_) ->
Fun = fun(#{aggregated := Aggr, per_client := Cli} = Bucket) ->
Aggr2 = Aggr#{rate := 1,
capacity := 200,
initial := 50},
Cli2 = Cli#{rate := infinity, capacity := infinity, divisible := true,
max_retry_time := 100, failure_strategy := force},
Bucket#{aggregated := Aggr2,
per_client := Cli2}
end,
Case = fun() ->
Client = connect(default),
{_, _, Retry, L2} = emqx_htb_limiter:check(150, Client),
timer:sleep(200),
{ok, L3} = emqx_htb_limiter:check(Retry, L2),
Avaiable = emqx_htb_limiter:available(L3),
?assert(Avaiable >= 50)
end,
with_bucket(default, Fun, Case).
t_short_board(_) ->
Fun = fun(#{aggregated := Aggr, per_client := Cli} = Bucket) ->
Aggr2 = Aggr#{rate := ?RATE("100/1s"),
@ -285,7 +356,7 @@ t_limit_zone_with_unlimit_bucket(_) ->
t_burst_and_fairness(_) ->
GlobalMod = fun(Cfg) ->
Cfg#{burst := ?RATE("60/1s")}
end,
end,
ZoneMod = fun(Cfg) ->
Cfg#{rate := ?RATE("600/1s"),
@ -319,9 +390,50 @@ t_burst_and_fairness(_) ->
[{b1, Bucket}, {b2, Bucket}],
Case).
t_burst(_) ->
GlobalMod = fun(Cfg) ->
Cfg#{burst := ?RATE("60/1s")}
end,
ZoneMod = fun(Cfg) ->
Cfg#{rate := ?RATE("60/1s"),
burst := ?RATE("60/1s")}
end,
Bucket = fun(#{aggregated := Aggr, per_client := Cli} = Bucket) ->
Aggr2 = Aggr#{rate := ?RATE("500/1s"),
initial := 0,
capacity := 500},
Cli2 = Cli#{rate := ?RATE("600/1s"),
capacity := 600,
divisible := true},
Bucket#{aggregated := Aggr2,
per_client := Cli2}
end,
Case = fun() ->
C1 = counters:new(1, []),
C2 = counters:new(1, []),
C3 = counters:new(1, []),
start_client(b1, ?NOW + 2000, C1, 20),
start_client(b2, ?NOW + 2000, C2, 30),
start_client(b3, ?NOW + 2000, C3, 30),
timer:sleep(2100),
Total = lists:sum([counters:get(X, 1) || X <- [C1, C2, C3]]),
in_range(Total / 2, 30)
end,
with_global(GlobalMod,
default,
ZoneMod,
[{b1, Bucket}, {b2, Bucket}, {b3, Bucket}],
Case).
t_limit_global_with_unlimit_other(_) ->
GlobalMod = fun(Cfg) ->
Cfg#{rate := ?RATE("600/1s")}
Cfg#{rate := ?RATE("600/1s")}
end,
ZoneMod = fun(Cfg) -> Cfg#{rate := infinity} end,
@ -483,14 +595,85 @@ t_zone_hunger_and_fair(_) ->
[{b1, Bucket(z1, ?RATE("600/1s"))}, {b2, Bucket(z2, ?RATE("50/1s"))}],
Case).
%%--------------------------------------------------------------------
%% Test Cases container
%%--------------------------------------------------------------------
t_new_container(_) ->
C1 = emqx_limiter_container:new(),
C2 = emqx_limiter_container:new([message_routing]),
C3 = emqx_limiter_container:update_by_name(message_routing, default, C1),
?assertMatch(#{message_routing := _,
retry_ctx := undefined,
{retry, message_routing} := _}, C2),
?assertMatch(#{message_routing := _,
retry_ctx := undefined,
{retry, message_routing} := _}, C3),
ok.
t_check_container(_) ->
Cfg = fun(Cfg) ->
Cfg#{rate := ?RATE("1000/1s"),
initial := 1000,
capacity := 1000}
end,
Case = fun() ->
C1 = emqx_limiter_container:new([message_routing]),
{ok, C2} = emqx_limiter_container:check(1000, message_routing, C1),
{pause, Pause, C3} = emqx_limiter_container:check(1000, message_routing, C2),
timer:sleep(Pause),
{ok, C4} = emqx_limiter_container:retry(message_routing, C3),
Context = test,
C5 = emqx_limiter_container:set_retry_context(Context, C4),
RetryData = emqx_limiter_container:get_retry_context(C5),
?assertEqual(Context, RetryData)
end,
with_per_client(default, Cfg, Case).
%%--------------------------------------------------------------------
%% Test Cases misc
%%--------------------------------------------------------------------
t_limiter_manager(_) ->
{error, _} = emqx_limiter_manager:start_server(message_routing),
ignore = gen_server:call(emqx_limiter_manager, unexpected_call),
ok = gen_server:cast(emqx_limiter_manager, unexpected_cast),
erlang:send(erlang:whereis(emqx_limiter_manager), unexpected_info),
ok = emqx_limiter_manager:format_status(normal, ok),
ok.
t_limiter_app(_) ->
try
_ = emqx_limiter_app:start(undefined, undefined)
catch _:_ ->
ok
end,
ok = emqx_limiter_app:stop(undefined),
ok.
t_limiter_server(_) ->
State = emqx_limiter_server:info(message_routing),
?assertMatch(#{root := _,
counter := _,
index := _,
zones := _,
buckets := _,
nodes := _,
type := message_routing}, State),
Name = emqx_limiter_server:name(message_routing),
ignored = gen_server:call(Name, unexpected_call),
ok = gen_server:cast(Name, unexpected_cast),
erlang:send(erlang:whereis(Name), unexpected_info),
ok = emqx_limiter_server:format_status(normal, ok),
ok.
%%--------------------------------------------------------------------
%%% Internal functions
%%--------------------------------------------------------------------
start_client(Name, EndTime, Counter, Number) ->
lists:foreach(fun(_) ->
spawn(fun() ->
start_client(Name, EndTime, Counter)
end)
spawn(fun() ->
start_client(Name, EndTime, Counter)
end)
end,
lists:seq(1, Number)).
@ -612,16 +795,18 @@ with_config(Path, Modifier, Case) ->
emqx_config:put(Path, NewCfg),
emqx_limiter_manager:restart_server(message_routing),
timer:sleep(100),
DelayReturn
= try
Return = Case(),
fun() -> Return end
catch Type:Reason:Trace ->
fun() -> erlang:raise(Type, Reason, Trace) end
end,
DelayReturn = delay_return(Case),
emqx_config:put(Path, Cfg),
DelayReturn().
delay_return(Case) ->
try
Return = Case(),
fun() -> Return end
catch Type:Reason:Trace ->
fun() -> erlang:raise(Type, Reason, Trace) end
end.
connect(Name) ->
emqx_limiter_server:connect(message_routing, Name).
@ -636,7 +821,7 @@ print_average_rate(Counter, Second) ->
PerSec = Cost / Second,
ct:pal("Cost:~p PerSec:~p ~n", [Cost, PerSec]).
in_range(Val, Expected) when Val < Expected * 0.6 ->
in_range(Val, Expected) when Val < Expected * 0.5 ->
ct:pal("Val:~p smaller than min bound", [Val]),
false;
in_range(Val, Expected) when Val > Expected * 1.8 ->
@ -663,6 +848,6 @@ apply_modifier(Name, Modifier, #{default := Template} = Cfg) ->
apply_modifier(Pairs, #{default := Template}) ->
Fun = fun({N, M}, Acc) ->
Acc#{N => M(Template)}
Acc#{N => M(Template)}
end,
lists:foldl(Fun, #{}, Pairs).