yuanbiao
/
emqx
1
0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
emqx/apps/emqx_modules/src/emqx_delayed.erl

491 lines
15 KiB
Erlang
Raw Permalink Blame History

%%--------------------------------------------------------------------
%% Copyright (c) 2020-2024 EMQ Technologies Co., Ltd. All Rights Reserved.
%%
%% Licensed under the Apache License, Version 2.0 (the "License");
%% you may not use this file except in compliance with the License.
%% You may obtain a copy of the License at
%%
%% http://www.apache.org/licenses/LICENSE-2.0
%%
%% Unless required by applicable law or agreed to in writing, software
%% distributed under the License is distributed on an "AS IS" BASIS,
%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
%% See the License for the specific language governing permissions and
%% limitations under the License.
%%--------------------------------------------------------------------
-module(emqx_delayed).
-behaviour(gen_server).
-behaviour(emqx_config_handler).
-include_lib("emqx/include/emqx.hrl").
-include_lib("emqx/include/types.hrl").
-include_lib("emqx/include/logger.hrl").
-include_lib("snabbkaffe/include/snabbkaffe.hrl").
-include_lib("emqx/include/emqx_hooks.hrl").
-export([
create_tables/0,
start_link/0,
on_message_publish/1
]).
%% gen_server callbacks
-export([
init/1,
handle_call/3,
handle_cast/2,
handle_info/2,
terminate/2,
code_change/3
]).
%% API
-export([
load/0,
unload/0,
load_or_unload/1,
get_conf/1,
update_config/1,
delayed_count/0,
list/1,
get_delayed_message/1,
get_delayed_message/2,
delete_delayed_message/1,
delete_delayed_message/2,
delete_delayed_messages_by_topic_name/1,
clear_all/0,
%% rpc target
clear_all_local/0,
cluster_list/1
]).
%% exports for internal rpc
-export([
do_delete_delayed_messages_by_topic_name/1
]).
%% exports for query
-export([
qs2ms/2,
format_delayed/1,
format_delayed/2
]).
-export([
post_config_update/5
]).
%% exported for `emqx_telemetry'
-export([get_basic_usage_info/0]).
-record(delayed_message, {key, delayed, msg}).
-type delayed_message() :: #delayed_message{}.
-type with_id_return() :: ok | {error, not_found}.
-type with_id_return(T) :: {ok, T} | {error, not_found}.
-export_type([with_id_return/0, with_id_return/1]).
-type state() :: #{
publish_timer := option(reference()),
publish_at := non_neg_integer(),
stats_timer := option(reference()),
stats_fun := option(fun((pos_integer()) -> ok))
}.
%% sync ms with record change
-define(QUERY_MS(Id), [{{delayed_message, {'_', Id}, '_', '_'}, [], ['$_']}]).
-define(DELETE_MS(Id), [{{delayed_message, {'$1', Id}, '_', '_'}, [], ['$1']}]).
-define(DELETE_BY_TOPIC_MS(Topic), [
{
{delayed_message, '$1', '_', {message, '_', '_', '_', '_', '_', Topic, '_', '_', '_'}},
[],
['$1']
}
]).
-define(TAB, ?MODULE).
-define(SERVER, ?MODULE).
-define(MAX_INTERVAL, 4294967).
-define(FORMAT_FUN, {?MODULE, format_delayed}).
-define(NOW, erlang:system_time(milli_seconds)).
%%------------------------------------------------------------------------------
%% Mnesia bootstrap
%%------------------------------------------------------------------------------
create_tables() ->
ok = mria:create_table(?TAB, [
{type, ordered_set},
{storage, disc_copies},
{local_content, true},
{record_name, delayed_message},
{attributes, record_info(fields, delayed_message)}
]),
[?TAB].
%%------------------------------------------------------------------------------
%% Hooks
%%------------------------------------------------------------------------------
on_message_publish(
Msg = #message{
id = Id,
topic = <<"$delayed/", Topic/binary>>,
timestamp = Ts
}
) ->
[Delay, Topic1] = binary:split(Topic, <<"/">>),
{PubAt, Delayed} =
case binary_to_integer(Delay) of
Interval when Interval < ?MAX_INTERVAL ->
{Interval * 1000 + Ts, Interval};
Timestamp ->
%% Check malicious timestamp?
Internal = Timestamp - erlang:round(Ts / 1000),
case Internal > ?MAX_INTERVAL of
true -> error(invalid_delayed_timestamp);
false -> {Timestamp * 1000, Internal}
end
end,
PubMsg = Msg#message{topic = Topic1},
Headers = PubMsg#message.headers,
case store(#delayed_message{key = {PubAt, Id}, delayed = Delayed, msg = PubMsg}) of
ok -> ok;
{error, Error} -> ?SLOG(error, #{msg => "store_delayed_message_fail", error => Error})
end,
{stop, PubMsg#message{headers = Headers#{allow_publish => false}}};
on_message_publish(Msg) ->
{ok, Msg}.
%%------------------------------------------------------------------------------
%% Start delayed publish server
%%------------------------------------------------------------------------------
-spec start_link() -> emqx_types:startlink_ret().
start_link() ->
gen_server:start_link({local, ?SERVER}, ?MODULE, [], []).
-spec store(delayed_message()) -> ok | {error, atom()}.
store(DelayedMsg) ->
gen_server:call(?SERVER, {store, DelayedMsg}, infinity).
get_conf(Key) ->
emqx_conf:get([delayed, Key]).
load() ->
load_or_unload(true).
unload() ->
load_or_unload(false).
load_or_unload(Bool) ->
gen_server:call(?SERVER, {do_load_or_unload, Bool}).
-spec delayed_count() -> non_neg_integer().
delayed_count() -> mnesia:table_info(?TAB, size).
list(Params) ->
emqx_mgmt_api:paginate(?TAB, Params, ?FORMAT_FUN).
cluster_list(Params) ->
emqx_mgmt_api:cluster_query(
?TAB,
Params,
[],
fun ?MODULE:qs2ms/2,
fun ?MODULE:format_delayed/2
).
-spec qs2ms(atom(), {list(), list()}) -> emqx_mgmt_api:match_spec_and_filter().
qs2ms(_Table, {_Qs, _Fuzzy}) ->
#{
match_spec => [{'$1', [], ['$1']}],
fuzzy_fun => undefined
}.
format_delayed(Delayed) ->
format_delayed(node(), Delayed).
format_delayed(WhichNode, Delayed) ->
format_delayed(WhichNode, Delayed, false).
format_delayed(
WhichNode,
#delayed_message{
key = {ExpectTimeStamp, Id},
delayed = Delayed,
msg = #message{
topic = Topic,
from = From,
headers = Headers,
qos = Qos,
timestamp = PublishTimeStamp,
payload = Payload
}
},
WithPayload
) ->
PublishTime = emqx_utils_calendar:epoch_to_rfc3339(PublishTimeStamp),
ExpectTime = emqx_utils_calendar:epoch_to_rfc3339(ExpectTimeStamp),
RemainingTime = ExpectTimeStamp - ?NOW,
Result = #{
msgid => emqx_guid:to_hexstr(Id),
node => WhichNode,
publish_at => PublishTime,
delayed_interval => Delayed,
delayed_remaining => RemainingTime div 1000,
expected_at => ExpectTime,
topic => Topic,
qos => Qos,
from_clientid => From,
from_username => maps:get(username, Headers, undefined)
},
case WithPayload of
true ->
Result#{payload => Payload};
_ ->
Result
end.
-spec get_delayed_message(binary()) -> with_id_return(map()).
get_delayed_message(Id) ->
case ets:select(?TAB, ?QUERY_MS(Id)) of
[] ->
{error, not_found};
Rows ->
Message = hd(Rows),
{ok, format_delayed(node(), Message, true)}
end.
get_delayed_message(Node, Id) when Node =:= node() ->
get_delayed_message(Id);
get_delayed_message(Node, Id) ->
emqx_delayed_proto_v2:get_delayed_message(Node, Id).
-spec delete_delayed_message(binary()) -> with_id_return().
delete_delayed_message(Id) ->
case ets:select(?TAB, ?DELETE_MS(Id)) of
[] ->
{error, not_found};
Rows ->
Timestamp = hd(Rows),
mria:dirty_delete(?TAB, {Timestamp, Id})
end.
delete_delayed_message(Node, Id) when Node =:= node() ->
delete_delayed_message(Id);
delete_delayed_message(Node, Id) ->
emqx_delayed_proto_v2:delete_delayed_message(Node, Id).
-spec delete_delayed_messages_by_topic_name(binary()) -> with_id_return().
delete_delayed_messages_by_topic_name(TopicName) when is_binary(TopicName) ->
Nodes = emqx:running_nodes(),
Result = emqx_delayed_proto_v3:delete_delayed_messages_by_topic_name(Nodes, TopicName),
case
lists:any(
fun
({ok, ok}) -> true;
(_) -> false
end,
Result
)
of
true ->
ok;
false ->
Errors = lists:filter(
fun
({ok, {error, not_found}}) -> false;
(_) -> true
end,
Result
),
case Errors of
[] ->
{error, not_found};
[Exception | _] ->
{error, Exception}
end
end.
-spec do_delete_delayed_messages_by_topic_name(binary()) -> with_id_return().
do_delete_delayed_messages_by_topic_name(TopicName) when is_binary(TopicName) ->
case ets:select(?TAB, ?DELETE_BY_TOPIC_MS(TopicName)) of
[] ->
{error, not_found};
Rows ->
lists:foreach(fun(Key) -> mria:dirty_delete(?TAB, Key) end, Rows)
end.
-spec clear_all() -> ok.
clear_all() ->
Nodes = emqx:running_nodes(),
_ = emqx_delayed_proto_v2:clear_all(Nodes),
ok.
%% rpc target
-spec clear_all_local() -> ok.
clear_all_local() ->
_ = mria:clear_table(?TAB),
ok.
update_config(Config) ->
emqx_conf:update([delayed], Config, #{rawconf_with_defaults => true, override_to => cluster}).
post_config_update(_KeyPath, _ConfigReq, NewConf, _OldConf, _AppEnvs) ->
Enable = maps:get(enable, NewConf, undefined),
load_or_unload(Enable).
%%------------------------------------------------------------------------------
%% gen_server callback
%%------------------------------------------------------------------------------
init([]) ->
ok = mria:wait_for_tables([?TAB]),
erlang:process_flag(trap_exit, true),
emqx_conf:add_handler([delayed], ?MODULE),
State =
ensure_stats_event(
ensure_publish_timer(#{
publish_timer => undefined,
publish_at => 0,
stats_timer => undefined,
stats_fun => undefined
})
),
{ok, do_load_or_unload(emqx:get_config([delayed, enable]), State)}.
handle_call({store, DelayedMsg = #delayed_message{key = Key}}, _From, State) ->
Size = mnesia:table_info(?TAB, size),
case get_conf(max_delayed_messages) of
0 ->
ok = mria:dirty_write(?TAB, DelayedMsg),
emqx_metrics:inc('messages.delayed'),
{reply, ok, ensure_publish_timer(Key, State)};
Max when Size >= Max ->
{reply, {error, max_delayed_messages_full}, State};
Max when Size < Max ->
ok = mria:dirty_write(?TAB, DelayedMsg),
emqx_metrics:inc('messages.delayed'),
{reply, ok, ensure_publish_timer(Key, State)}
end;
handle_call({do_load_or_unload, Bool}, _From, State0) ->
State = do_load_or_unload(Bool, State0),
{reply, ok, State};
handle_call(Req, _From, State) ->
?tp(error, emqx_delayed_unexpected_call, #{call => Req}),
{reply, ignored, State}.
handle_cast(Msg, State) ->
?tp(error, emqx_delayed_unexpected_cast, #{cast => Msg}),
{noreply, State}.
%% Do Publish...
handle_info({timeout, TRef, do_publish}, State = #{publish_timer := TRef}) ->
DeletedKeys = do_publish(mnesia:dirty_first(?TAB), ?NOW),
lists:foreach(fun(Key) -> mria:dirty_delete(?TAB, Key) end, DeletedKeys),
{noreply, ensure_publish_timer(State#{publish_timer := undefined, publish_at := 0})};
handle_info(stats, State = #{stats_fun := StatsFun}) ->
StatsTimer = erlang:send_after(timer:seconds(1), self(), stats),
StatsFun(delayed_count()),
{noreply, State#{stats_timer := StatsTimer}, hibernate};
handle_info(Info, State) ->
?tp(error, emqx_delayed_unexpected_info, #{info => Info}),
{noreply, State}.
terminate(_Reason, #{stats_timer := StatsTimer} = State) ->
emqx_conf:remove_handler([delayed]),
emqx_utils:cancel_timer(StatsTimer),
do_load_or_unload(false, State).
code_change(_Vsn, State, _Extra) ->
{ok, State}.
%%------------------------------------------------------------------------------
%% Telemetry
%%------------------------------------------------------------------------------
-spec get_basic_usage_info() -> #{delayed_message_count => non_neg_integer()}.
get_basic_usage_info() ->
DelayedCount =
case ets:info(?TAB, size) of
undefined -> 0;
Num -> Num
end,
#{delayed_message_count => DelayedCount}.
%%------------------------------------------------------------------------------
%% Internal functions
%%------------------------------------------------------------------------------
%% Ensure the stats
-spec ensure_stats_event(state()) -> state().
ensure_stats_event(State) ->
StatsFun = emqx_stats:statsfun('delayed.count', 'delayed.max'),
StatsTimer = erlang:send_after(timer:seconds(1), self(), stats),
State#{stats_fun := StatsFun, stats_timer := StatsTimer}.
%% Ensure publish timer
-spec ensure_publish_timer(state()) -> state().
ensure_publish_timer(State) ->
ensure_publish_timer(mnesia:dirty_first(?TAB), State).
ensure_publish_timer('$end_of_table', State) ->
State#{publish_timer := undefined, publish_at := 0};
ensure_publish_timer({Ts, _Id}, State = #{publish_timer := undefined}) ->
ensure_publish_timer(Ts, ?NOW, State);
ensure_publish_timer({Ts, _Id}, State = #{publish_timer := TRef, publish_at := PubAt}) when
Ts < PubAt
->
ok = emqx_utils:cancel_timer(TRef),
ensure_publish_timer(Ts, ?NOW, State);
ensure_publish_timer(_Key, State) ->
State.
ensure_publish_timer(Ts, Now, State) ->
Interval = max(1, Ts - Now),
TRef = emqx_utils:start_timer(Interval, do_publish),
State#{publish_timer := TRef, publish_at := Now + Interval}.
do_publish(Key, Now) ->
do_publish(Key, Now, []).
%% Do publish
do_publish('$end_of_table', _Now, Acc) ->
Acc;
do_publish({Ts, _Id}, Now, Acc) when Ts > Now ->
Acc;
do_publish(Key = {Ts, _Id}, Now, Acc) when Ts =< Now ->
case mnesia:dirty_read(?TAB, Key) of
[] ->
ok;
[#delayed_message{msg = Msg}] ->
case emqx_banned:look_up({clientid, Msg#message.from}) of
[] ->
emqx_pool:async_submit(fun emqx:publish/1, [Msg]);
_ ->
?tp(
notice,
ignore_delayed_message_publish,
#{
reason => "client is banned",
clientid => Msg#message.from
}
),
ok
end
end,
do_publish(mnesia:dirty_next(?TAB, Key), Now, [Key | Acc]).
do_load_or_unload(true, State) ->
emqx_hooks:put('message.publish', {?MODULE, on_message_publish, []}, ?HP_DELAY_PUB),
State;
do_load_or_unload(false, #{publish_timer := PubTimer} = State) ->
emqx_hooks:del('message.publish', {?MODULE, on_message_publish}),
emqx_utils:cancel_timer(PubTimer),
ets:delete_all_objects(?TAB),
State#{publish_timer := undefined, publish_at := 0};
do_load_or_unload(_, State) ->
State.
Reference in New Issue View Git Blame Copy Permalink