emqx/apps/emqx_modules/src/emqx_delayed.erl

308 lines
10 KiB
Erlang

%%--------------------------------------------------------------------
%% Copyright (c) 2020-2021 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).
-include_lib("emqx/include/emqx.hrl").
-include_lib("emqx/include/logger.hrl").
%% Mnesia bootstrap
-export([mnesia/1]).
-boot_mnesia({mnesia, [boot]}).
-export([ 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
]).
%% gen_server callbacks
-export([ enable/0
, disable/0
, set_max_delayed_messages/1
, update_config/1
, list/1
, get_delayed_message/1
, delete_delayed_message/1
]).
-export([format_delayed/1]).
-record(delayed_message, {key, delayed, msg}).
%% sync ms with record change
-define(QUERY_MS(Id), [{{delayed_message, {'_', Id}, '_', '_'}, [], ['$_']}]).
-define(DELETE_MS(Id), [{{delayed_message, {'$1', Id}, '_', '_'}, [], ['$1']}]).
-define(TAB, ?MODULE).
-define(SERVER, ?MODULE).
-define(MAX_INTERVAL, 4294967).
-define(FORMAT_FUN, {?MODULE, format_delayed}).
%%--------------------------------------------------------------------
%% Mnesia bootstrap
%%--------------------------------------------------------------------
mnesia(boot) ->
ok = mria:create_table(?TAB, [
{type, ordered_set},
{storage, disc_copies},
{local_content, true},
{record_name, delayed_message},
{attributes, record_info(fields, delayed_message)}]).
%%--------------------------------------------------------------------
%% 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 + erlang:round(Ts / 1000), Interval};
Timestamp ->
%% Check malicious timestamp?
case (Timestamp - erlang:round(Ts / 1000)) > ?MAX_INTERVAL of
true -> error(invalid_delayed_timestamp);
false -> {Timestamp, Timestamp - erlang:round(Ts / 1000)}
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() ->
Opts = emqx_conf:get([delayed], #{}),
gen_server:start_link({local, ?SERVER}, ?MODULE, [Opts], []).
-spec(store(#delayed_message{}) -> ok | {error, atom()}).
store(DelayedMsg) ->
gen_server:call(?SERVER, {store, DelayedMsg}, infinity).
enable() ->
gen_server:call(?SERVER, enable).
disable() ->
gen_server:call(?SERVER, disable).
set_max_delayed_messages(Max) ->
gen_server:call(?SERVER, {set_max_delayed_messages, Max}).
list(Params) ->
emqx_mgmt_api:paginate(?TAB, Params, ?FORMAT_FUN).
format_delayed(Delayed) ->
format_delayed(Delayed, false).
format_delayed(#delayed_message{key = {ExpectTimeStamp, Id}, delayed = Delayed,
msg = #message{topic = Topic,
from = From,
headers = Headers,
qos = Qos,
timestamp = PublishTimeStamp,
payload = Payload}}, WithPayload) ->
PublishTime = to_rfc3339(PublishTimeStamp div 1000),
ExpectTime = to_rfc3339(ExpectTimeStamp),
RemainingTime = ExpectTimeStamp - erlang:system_time(second),
Result = #{
msgid => emqx_guid:to_hexstr(Id),
publish_at => PublishTime,
delayed_interval => Delayed,
delayed_remaining => RemainingTime,
expected_at => ExpectTime,
topic => Topic,
qos => Qos,
from_clientid => From,
from_username => maps:get(username, Headers, undefined)
},
case WithPayload of
true ->
Result#{payload => base64:encode(Payload)};
_ ->
Result
end.
to_rfc3339(Timestamp) ->
list_to_binary(calendar:system_time_to_rfc3339(Timestamp, [{unit, second}])).
get_delayed_message(Id0) ->
try emqx_guid:from_hexstr(Id0) of
Id ->
case ets:select(?TAB, ?QUERY_MS(Id)) of
[] ->
{error, not_found};
Rows ->
Message = hd(Rows),
{ok, format_delayed(Message, true)}
end
catch
error:function_clause -> {error, id_schema_error}
end.
delete_delayed_message(Id0) ->
Id = emqx_guid:from_hexstr(Id0),
case ets:select(?TAB, ?DELETE_MS(Id)) of
[] ->
{error, not_found};
Rows ->
Timestamp = hd(Rows),
mria:dirty_delete(?TAB, {Timestamp, Id})
end.
update_config(Config) ->
{ok, _} = emqx:update_config([delayed], Config).
%%--------------------------------------------------------------------
%% gen_server callback
%%--------------------------------------------------------------------
init([Opts]) ->
MaxDelayedMessages = maps:get(max_delayed_messages, Opts, 0),
{ok, ensure_stats_event(
ensure_publish_timer(#{timer => undefined,
publish_at => 0,
max_delayed_messages => MaxDelayedMessages}))}.
handle_call({set_max_delayed_messages, Max}, _From, State) ->
{reply, ok, State#{max_delayed_messages => Max}};
handle_call({store, DelayedMsg = #delayed_message{key = Key}},
_From, State = #{max_delayed_messages := 0}) ->
ok = mria:dirty_write(?TAB, DelayedMsg),
emqx_metrics:inc('messages.delayed'),
{reply, ok, ensure_publish_timer(Key, State)};
handle_call({store, DelayedMsg = #delayed_message{key = Key}},
_From, State = #{max_delayed_messages := Max}) ->
Size = mnesia:table_info(?TAB, size),
case Size >= Max of
true ->
{reply, {error, max_delayed_messages_full}, State};
false ->
ok = mria:dirty_write(?TAB, DelayedMsg),
emqx_metrics:inc('messages.delayed'),
{reply, ok, ensure_publish_timer(Key, State)}
end;
handle_call(enable, _From, State) ->
emqx_hooks:put('message.publish', {?MODULE, on_message_publish, []}),
{reply, ok, State};
handle_call(disable, _From, State) ->
emqx_hooks:del('message.publish', {?MODULE, on_message_publish}),
{reply, ok, State};
handle_call(Req, _From, State) ->
?SLOG(error, #{msg => "unexpected_call", call => Req}),
{reply, ignored, State}.
handle_cast(Msg, State) ->
?SLOG(error, #{msg => "unexpected_cast", cast => Msg}),
{noreply, State}.
%% Do Publish...
handle_info({timeout, TRef, do_publish}, State = #{timer := TRef}) ->
DeletedKeys = do_publish(mnesia:dirty_first(?TAB), os:system_time(seconds)),
lists:foreach(fun(Key) -> mria:dirty_delete(?TAB, Key) end, DeletedKeys),
{noreply, ensure_publish_timer(State#{timer := undefined, publish_at := 0})};
handle_info(stats, State = #{stats_fun := StatsFun}) ->
StatsFun(delayed_count()),
{noreply, State, hibernate};
handle_info(Info, State) ->
?SLOG(error, #{msg => "unexpected_info", info => Info}),
{noreply, State}.
terminate(_Reason, #{timer := TRef}) ->
emqx_misc:cancel_timer(TRef).
code_change(_Vsn, State, _Extra) ->
{ok, State}.
%%--------------------------------------------------------------------
%% Internal functions
%%--------------------------------------------------------------------
%% Ensure the stats
ensure_stats_event(State) ->
StatsFun = emqx_stats:statsfun('delayed.count', 'delayed.max'),
{ok, StatsTimer} = timer:send_interval(timer:seconds(1), stats),
State#{stats_fun => StatsFun, stats_timer => StatsTimer}.
%% Ensure publish timer
ensure_publish_timer(State) ->
ensure_publish_timer(mnesia:dirty_first(?TAB), State).
ensure_publish_timer('$end_of_table', State) ->
State#{timer := undefined, publish_at := 0};
ensure_publish_timer({Ts, _Id}, State = #{timer := undefined}) ->
ensure_publish_timer(Ts, os:system_time(seconds), State);
ensure_publish_timer({Ts, _Id}, State = #{timer := TRef, publish_at := PubAt})
when Ts < PubAt ->
ok = emqx_misc:cancel_timer(TRef),
ensure_publish_timer(Ts, os:system_time(seconds), State);
ensure_publish_timer(_Key, State) ->
State.
ensure_publish_timer(Ts, Now, State) ->
Interval = max(1, Ts - Now),
TRef = emqx_misc:start_timer(timer:seconds(Interval), do_publish),
State#{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}] ->
emqx_pool:async_submit(fun emqx:publish/1, [Msg])
end,
do_publish(mnesia:dirty_next(?TAB, Key), Now, [Key|Acc]).
-spec(delayed_count() -> non_neg_integer()).
delayed_count() -> mnesia:table_info(?TAB, size).