feat: throttle with resource_id

This commit is contained in:
zhongwencool 2024-07-25 13:06:35 +08:00 committed by zmstone
parent 2924ec582a
commit f6f1d32da0
5 changed files with 149 additions and 46 deletions

View File

@ -38,16 +38,20 @@
). ).
%% NOTE: do not forget to use atom for msg and add every used msg to %% NOTE: do not forget to use atom for msg and add every used msg to
%% the default value of `log.thorttling.msgs` list. %% the default value of `log.throttling.msgs` list.
-define(SLOG_THROTTLE(Level, Data), -define(SLOG_THROTTLE(Level, Data),
?SLOG_THROTTLE(Level, Data, #{}) ?SLOG_THROTTLE(Level, Data, #{})
). ).
-define(SLOG_THROTTLE(Level, Data, Meta), -define(SLOG_THROTTLE(Level, Data, Meta),
?SLOG_THROTTLE(Level, undefined, Data, Meta)
).
-define(SLOG_THROTTLE(Level, UniqueKey, Data, Meta),
case logger:allow(Level, ?MODULE) of case logger:allow(Level, ?MODULE) of
true -> true ->
(fun(#{msg := __Msg} = __Data) -> (fun(#{msg := __Msg} = __Data) ->
case emqx_log_throttler:allow(__Msg) of case emqx_log_throttler:allow(__Msg, UniqueKey) of
true -> true ->
logger:log(Level, __Data, Meta); logger:log(Level, __Data, Meta);
false -> false ->

View File

@ -25,7 +25,7 @@
-export([start_link/0]). -export([start_link/0]).
%% throttler API %% throttler API
-export([allow/1]). -export([allow/2]).
%% gen_server callbacks %% gen_server callbacks
-export([ -export([
@ -40,23 +40,22 @@
-define(SEQ_ID(Msg), {?MODULE, Msg}). -define(SEQ_ID(Msg), {?MODULE, Msg}).
-define(NEW_SEQ, atomics:new(1, [{signed, false}])). -define(NEW_SEQ, atomics:new(1, [{signed, false}])).
-define(GET_SEQ(Msg), persistent_term:get(?SEQ_ID(Msg), undefined)). -define(GET_SEQ(Msg), persistent_term:get(?SEQ_ID(Msg), undefined)).
-define(ERASE_SEQ(Msg), persistent_term:erase(?SEQ_ID(Msg))).
-define(RESET_SEQ(SeqRef), atomics:put(SeqRef, 1, 0)). -define(RESET_SEQ(SeqRef), atomics:put(SeqRef, 1, 0)).
-define(INC_SEQ(SeqRef), atomics:add(SeqRef, 1, 1)). -define(INC_SEQ(SeqRef), atomics:add(SeqRef, 1, 1)).
-define(GET_DROPPED(SeqRef), atomics:get(SeqRef, 1) - 1). -define(GET_DROPPED(SeqRef), atomics:get(SeqRef, 1) - 1).
-define(IS_ALLOWED(SeqRef), atomics:add_get(SeqRef, 1, 1) =:= 1). -define(IS_ALLOWED(SeqRef), atomics:add_get(SeqRef, 1, 1) =:= 1).
-define(NEW_THROTTLE(Msg, SeqRef), persistent_term:put(?SEQ_ID(Msg), SeqRef)).
-define(MSGS_LIST, emqx:get_config([log, throttling, msgs], [])). -define(MSGS_LIST, emqx:get_config([log, throttling, msgs], [])).
-define(TIME_WINDOW_MS, timer:seconds(emqx:get_config([log, throttling, time_window], 60))). -define(TIME_WINDOW_MS, timer:seconds(emqx:get_config([log, throttling, time_window], 60))).
-spec allow(atom()) -> boolean(). -spec allow(atom(), any()) -> boolean().
allow(Msg) when is_atom(Msg) -> allow(Msg, UniqueKey) when is_atom(Msg) ->
case emqx_logger:get_primary_log_level() of case emqx_logger:get_primary_log_level() of
debug -> debug ->
true; true;
_ -> _ ->
do_allow(Msg) do_allow(Msg, UniqueKey)
end. end.
-spec start_link() -> startlink_ret(). -spec start_link() -> startlink_ret().
@ -68,7 +67,8 @@ start_link() ->
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
init([]) -> init([]) ->
ok = lists:foreach(fun(Msg) -> ?NEW_THROTTLE(Msg, ?NEW_SEQ) end, ?MSGS_LIST), process_flag(trap_exit, true),
ok = lists:foreach(fun(Msg) -> new_throttler(Msg) end, ?MSGS_LIST),
CurrentPeriodMs = ?TIME_WINDOW_MS, CurrentPeriodMs = ?TIME_WINDOW_MS,
TimerRef = schedule_refresh(CurrentPeriodMs), TimerRef = schedule_refresh(CurrentPeriodMs),
{ok, #{timer_ref => TimerRef, current_period_ms => CurrentPeriodMs}}. {ok, #{timer_ref => TimerRef, current_period_ms => CurrentPeriodMs}}.
@ -88,14 +88,19 @@ handle_info(refresh, #{current_period_ms := PeriodMs} = State) ->
case ?GET_SEQ(Msg) of case ?GET_SEQ(Msg) of
%% Should not happen, unless the static ids list is updated at run-time. %% Should not happen, unless the static ids list is updated at run-time.
undefined -> undefined ->
?NEW_THROTTLE(Msg, ?NEW_SEQ), new_throttler(Msg),
?tp(log_throttler_new_msg, #{throttled_msg => Msg}), ?tp(log_throttler_new_msg, #{throttled_msg => Msg}),
Acc; Acc;
SeqMap when is_map(SeqMap) ->
maps:fold(
fun(Key, Ref, Acc0) ->
drop_stats(Ref, emqx_utils:format("~ts:~s", [Msg, Key]), Acc0)
end,
Acc,
SeqMap
);
SeqRef -> SeqRef ->
Dropped = ?GET_DROPPED(SeqRef), drop_stats(SeqRef, Msg, Acc)
ok = ?RESET_SEQ(SeqRef),
?tp(log_throttler_dropped, #{dropped_count => Dropped, throttled_msg => Msg}),
maybe_add_dropped(Msg, Dropped, Acc)
end end
end, end,
#{}, #{},
@ -112,7 +117,34 @@ handle_info(Info, State) ->
?SLOG(error, #{msg => "unxpected_info", info => Info}), ?SLOG(error, #{msg => "unxpected_info", info => Info}),
{noreply, State}. {noreply, State}.
drop_stats(SeqRef, Msg, Acc) ->
Dropped = ?GET_DROPPED(SeqRef),
ok = ?RESET_SEQ(SeqRef),
?tp(log_throttler_dropped, #{dropped_count => Dropped, throttled_msg => Msg}),
maybe_add_dropped(Msg, Dropped, Acc).
terminate(_Reason, _State) -> terminate(_Reason, _State) ->
lists:foreach(
fun(Msg) ->
case ?GET_SEQ(Msg) of
undefined ->
ok;
SeqMap when is_map(SeqMap) ->
maps:foreach(
fun(_, Ref) ->
ok = ?RESET_SEQ(Ref)
end,
SeqMap
);
SeqRef ->
%% atomics don't have erase API...
%% (if nobody hold the ref, the atomics should erase automatically?)
ok = ?RESET_SEQ(SeqRef)
end,
?ERASE_SEQ(Msg)
end,
?MSGS_LIST
),
ok. ok.
code_change(_OldVsn, State, _Extra) -> code_change(_OldVsn, State, _Extra) ->
@ -122,17 +154,27 @@ code_change(_OldVsn, State, _Extra) ->
%% internal functions %% internal functions
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
do_allow(Msg) -> do_allow(Msg, UniqueKey) ->
case persistent_term:get(?SEQ_ID(Msg), undefined) of case persistent_term:get(?SEQ_ID(Msg), undefined) of
undefined -> undefined ->
%% This is either a race condition (emqx_log_throttler is not started yet) %% This is either a race condition (emqx_log_throttler is not started yet)
%% or a developer mistake (msg used in ?SLOG_THROTTLE/2,3 macro is %% or a developer mistake (msg used in ?SLOG_THROTTLE/2,3 macro is
%% not added to the default value of `log.throttling.msgs`. %% not added to the default value of `log.throttling.msgs`.
?SLOG(info, #{ ?SLOG(debug, #{
msg => "missing_log_throttle_sequence", msg => "log_throttle_disabled",
throttled_msg => Msg throttled_msg => Msg
}), }),
true; true;
%% e.g: unrecoverable msg throttle according resource_id
SeqMap when is_map(SeqMap) ->
case maps:find(UniqueKey, SeqMap) of
{ok, SeqRef} ->
?IS_ALLOWED(SeqRef);
error ->
SeqRef = ?NEW_SEQ,
new_throttler(Msg, SeqMap#{UniqueKey => SeqRef}),
true
end;
SeqRef -> SeqRef ->
?IS_ALLOWED(SeqRef) ?IS_ALLOWED(SeqRef)
end. end.
@ -154,3 +196,11 @@ maybe_log_dropped(_DroppedStats, _PeriodMs) ->
schedule_refresh(PeriodMs) -> schedule_refresh(PeriodMs) ->
?tp(log_throttler_sched_refresh, #{new_period_ms => PeriodMs}), ?tp(log_throttler_sched_refresh, #{new_period_ms => PeriodMs}),
erlang:send_after(PeriodMs, ?MODULE, refresh). erlang:send_after(PeriodMs, ?MODULE, refresh).
new_throttler(unrecoverable_resource_error = Msg) ->
persistent_term:put(?SEQ_ID(Msg), #{});
new_throttler(Msg) ->
persistent_term:put(?SEQ_ID(Msg), ?NEW_SEQ).
new_throttler(Msg, Map) ->
persistent_term:put(?SEQ_ID(Msg), Map).

View File

@ -26,6 +26,7 @@
%% Have to use real msgs, as the schema is guarded by enum. %% Have to use real msgs, as the schema is guarded by enum.
-define(THROTTLE_MSG, authorization_permission_denied). -define(THROTTLE_MSG, authorization_permission_denied).
-define(THROTTLE_MSG1, cannot_publish_to_topic_due_to_not_authorized). -define(THROTTLE_MSG1, cannot_publish_to_topic_due_to_not_authorized).
-define(THROTTLE_UNRECOVERABLE_MSG, unrecoverable_resource_error).
-define(TIME_WINDOW, <<"1s">>). -define(TIME_WINDOW, <<"1s">>).
all() -> emqx_common_test_helpers:all(?MODULE). all() -> emqx_common_test_helpers:all(?MODULE).
@ -59,6 +60,11 @@ end_per_suite(Config) ->
emqx_cth_suite:stop(?config(suite_apps, Config)), emqx_cth_suite:stop(?config(suite_apps, Config)),
emqx_config:delete_override_conf_files(). emqx_config:delete_override_conf_files().
init_per_testcase(t_throttle_recoverable_msg, Config) ->
ok = snabbkaffe:start_trace(),
[?THROTTLE_MSG] = Conf = emqx:get_config([log, throttling, msgs]),
{ok, _} = emqx_conf:update([log, throttling, msgs], [?THROTTLE_UNRECOVERABLE_MSG | Conf], #{}),
Config;
init_per_testcase(t_throttle_add_new_msg, Config) -> init_per_testcase(t_throttle_add_new_msg, Config) ->
ok = snabbkaffe:start_trace(), ok = snabbkaffe:start_trace(),
[?THROTTLE_MSG] = Conf = emqx:get_config([log, throttling, msgs]), [?THROTTLE_MSG] = Conf = emqx:get_config([log, throttling, msgs]),
@ -72,6 +78,10 @@ init_per_testcase(_TC, Config) ->
ok = snabbkaffe:start_trace(), ok = snabbkaffe:start_trace(),
Config. Config.
end_per_testcase(t_throttle_recoverable_msg, _Config) ->
ok = snabbkaffe:stop(),
{ok, _} = emqx_conf:update([log, throttling, msgs], [?THROTTLE_MSG], #{}),
ok;
end_per_testcase(t_throttle_add_new_msg, _Config) -> end_per_testcase(t_throttle_add_new_msg, _Config) ->
ok = snabbkaffe:stop(), ok = snabbkaffe:stop(),
{ok, _} = emqx_conf:update([log, throttling, msgs], [?THROTTLE_MSG], #{}), {ok, _} = emqx_conf:update([log, throttling, msgs], [?THROTTLE_MSG], #{}),
@ -101,8 +111,8 @@ t_throttle(_Config) ->
5000 5000
), ),
?assert(emqx_log_throttler:allow(?THROTTLE_MSG)), ?assert(emqx_log_throttler:allow(?THROTTLE_MSG, undefined)),
?assertNot(emqx_log_throttler:allow(?THROTTLE_MSG)), ?assertNot(emqx_log_throttler:allow(?THROTTLE_MSG, undefined)),
{ok, _} = ?block_until( {ok, _} = ?block_until(
#{ #{
?snk_kind := log_throttler_dropped, ?snk_kind := log_throttler_dropped,
@ -115,14 +125,48 @@ t_throttle(_Config) ->
[] []
). ).
t_throttle_recoverable_msg(_Config) ->
ResourceId = <<"resource_id">>,
ThrottledMsg = emqx_utils:format("~ts:~s", [?THROTTLE_UNRECOVERABLE_MSG, ResourceId]),
?check_trace(
begin
%% Warm-up and block to increase the probability that next events
%% will be in the same throttling time window.
{ok, _} = ?block_until(
#{?snk_kind := log_throttler_new_msg, throttled_msg := ?THROTTLE_UNRECOVERABLE_MSG},
5000
),
{_, {ok, _}} = ?wait_async_action(
events(?THROTTLE_UNRECOVERABLE_MSG, ResourceId),
#{
?snk_kind := log_throttler_dropped,
throttled_msg := ThrottledMsg
},
5000
),
?assert(emqx_log_throttler:allow(?THROTTLE_UNRECOVERABLE_MSG, ResourceId)),
?assertNot(emqx_log_throttler:allow(?THROTTLE_UNRECOVERABLE_MSG, ResourceId)),
{ok, _} = ?block_until(
#{
?snk_kind := log_throttler_dropped,
throttled_msg := ThrottledMsg,
dropped_count := 1
},
3000
)
end,
[]
).
t_throttle_add_new_msg(_Config) -> t_throttle_add_new_msg(_Config) ->
?check_trace( ?check_trace(
begin begin
{ok, _} = ?block_until( {ok, _} = ?block_until(
#{?snk_kind := log_throttler_new_msg, throttled_msg := ?THROTTLE_MSG1}, 5000 #{?snk_kind := log_throttler_new_msg, throttled_msg := ?THROTTLE_MSG1}, 5000
), ),
?assert(emqx_log_throttler:allow(?THROTTLE_MSG1)), ?assert(emqx_log_throttler:allow(?THROTTLE_MSG1, undefined)),
?assertNot(emqx_log_throttler:allow(?THROTTLE_MSG1)), ?assertNot(emqx_log_throttler:allow(?THROTTLE_MSG1, undefined)),
{ok, _} = ?block_until( {ok, _} = ?block_until(
#{ #{
?snk_kind := log_throttler_dropped, ?snk_kind := log_throttler_dropped,
@ -137,8 +181,8 @@ t_throttle_add_new_msg(_Config) ->
t_throttle_no_msg(_Config) -> t_throttle_no_msg(_Config) ->
%% Must simply pass with no crashes %% Must simply pass with no crashes
?assert(emqx_log_throttler:allow(no_test_throttle_msg)), ?assert(emqx_log_throttler:allow(no_test_throttle_msg, undefined)),
?assert(emqx_log_throttler:allow(no_test_throttle_msg)), ?assert(emqx_log_throttler:allow(no_test_throttle_msg, undefined)),
timer:sleep(10), timer:sleep(10),
?assert(erlang:is_process_alive(erlang:whereis(emqx_log_throttler))). ?assert(erlang:is_process_alive(erlang:whereis(emqx_log_throttler))).
@ -168,8 +212,8 @@ t_throttle_debug_primary_level(_Config) ->
#{?snk_kind := log_throttler_dropped, throttled_msg := ?THROTTLE_MSG}, #{?snk_kind := log_throttler_dropped, throttled_msg := ?THROTTLE_MSG},
5000 5000
), ),
?assert(emqx_log_throttler:allow(?THROTTLE_MSG)), ?assert(emqx_log_throttler:allow(?THROTTLE_MSG, undefined)),
?assertNot(emqx_log_throttler:allow(?THROTTLE_MSG)), ?assertNot(emqx_log_throttler:allow(?THROTTLE_MSG, undefined)),
{ok, _} = ?block_until( {ok, _} = ?block_until(
#{ #{
?snk_kind := log_throttler_dropped, ?snk_kind := log_throttler_dropped,
@ -187,10 +231,13 @@ t_throttle_debug_primary_level(_Config) ->
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
events(Msg) -> events(Msg) ->
events(100, Msg). events(100, Msg, undefined).
events(N, Msg) -> events(Msg, Id) ->
[emqx_log_throttler:allow(Msg) || _ <- lists:seq(1, N)]. events(100, Msg, Id).
events(N, Msg, Id) ->
[emqx_log_throttler:allow(Msg, Id) || _ <- lists:seq(1, N)].
module_exists(Mod) -> module_exists(Mod) ->
case erlang:module_loaded(Mod) of case erlang:module_loaded(Mod) of

View File

@ -167,12 +167,4 @@
). ).
-define(TAG, "RESOURCE"). -define(TAG, "RESOURCE").
-define(LOG_LEVEL(_L_),
case _L_ of
true -> info;
false -> warning
end
).
-define(TAG, "RESOURCE").
-define(RESOURCE_ALLOCATION_TAB, emqx_resource_allocations). -define(RESOURCE_ALLOCATION_TAB, emqx_resource_allocations).

View File

@ -981,11 +981,16 @@ handle_query_result_pure(Id, {error, Reason} = Error, HasBeenSent, TraceCTX) ->
true -> true ->
PostFn = PostFn =
fun() -> fun() ->
?SLOG_THROTTLE(error, #{ ?SLOG_THROTTLE(
resource_id => Id, error,
msg => unrecoverable_resource_error, Id,
reason => Reason #{
}), resource_id => Id,
msg => unrecoverable_resource_error,
reason => Reason
},
#{tag => ?TAG}
),
ok ok
end, end,
Counters = Counters =
@ -1025,11 +1030,16 @@ handle_query_async_result_pure(Id, {error, Reason} = Error, HasBeenSent, TraceCT
true -> true ->
PostFn = PostFn =
fun() -> fun() ->
?SLOG_THROTTLE(error, #{ ?SLOG_THROTTLE(
resource_id => Id, error,
msg => unrecoverable_resource_error, Id,
reason => Reason #{
}), resource_id => Id,
msg => unrecoverable_resource_error,
reason => Reason
},
#{tag => ?TAG}
),
ok ok
end, end,
Counters = Counters =