|
|
@@ -0,0 +1,195 @@
|
|
|
+%%--------------------------------------------------------------------
|
|
|
+%% Copyright (c) 2020 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_mod_delayed).
|
|
|
+
|
|
|
+-behaviour(gen_server).
|
|
|
+
|
|
|
+-include_lib("emqx/include/emqx.hrl").
|
|
|
+-include_lib("emqx/include/logger.hrl").
|
|
|
+
|
|
|
+%% emqx_gen_mod callbacks
|
|
|
+-export([ load/1
|
|
|
+ , unload/1
|
|
|
+ ]).
|
|
|
+
|
|
|
+-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
|
|
|
+ ]).
|
|
|
+
|
|
|
+-record(delayed_message,
|
|
|
+ { key
|
|
|
+ , msg
|
|
|
+ }).
|
|
|
+
|
|
|
+-define(TAB, ?MODULE).
|
|
|
+-define(SERVER, ?MODULE).
|
|
|
+-define(MAX_INTERVAL, 4294967).
|
|
|
+
|
|
|
+%%--------------------------------------------------------------------
|
|
|
+%% Load/Unload
|
|
|
+%%--------------------------------------------------------------------
|
|
|
+
|
|
|
+-spec(load(list()) -> ok).
|
|
|
+load(_Env) ->
|
|
|
+ emqx_mod_sup:start_child(?MODULE, worker),
|
|
|
+ emqx:hook('message.publish', {?MODULE, on_message_publish, []}).
|
|
|
+
|
|
|
+-spec(unload(list()) -> ok).
|
|
|
+unload(_Env) ->
|
|
|
+ emqx_mod_sup:stop_child(?MODULE),
|
|
|
+ emqx:unhook('message.publish', {?MODULE, on_message_publish, []}).
|
|
|
+
|
|
|
+%%--------------------------------------------------------------------
|
|
|
+%% Hooks
|
|
|
+%%--------------------------------------------------------------------
|
|
|
+
|
|
|
+on_message_publish(Msg = #message{id = Id, topic = <<"$delayed/", Topic/binary>>, timestamp = Ts}) ->
|
|
|
+ [Delay, Topic1] = binary:split(Topic, <<"/">>),
|
|
|
+ PubAt = case binary_to_integer(Delay) of
|
|
|
+ Interval when Interval < ?MAX_INTERVAL ->
|
|
|
+ Interval + erlang:round(Ts / 1000);
|
|
|
+ Timestamp ->
|
|
|
+ %% Check malicious timestamp?
|
|
|
+ case (Timestamp - erlang:round(Ts / 1000)) > ?MAX_INTERVAL of
|
|
|
+ true -> error(invalid_delayed_timestamp);
|
|
|
+ false -> Timestamp
|
|
|
+ end
|
|
|
+ end,
|
|
|
+ PubMsg = Msg#message{topic = Topic1},
|
|
|
+ Headers = case PubMsg#message.headers of
|
|
|
+ undefined -> #{};
|
|
|
+ Headers0 -> Headers0
|
|
|
+ end,
|
|
|
+ ok = store(#delayed_message{key = {PubAt, delayed_mid(Id)}, msg = PubMsg}),
|
|
|
+ {stop, PubMsg#message{headers = Headers#{allow_publish => false}}};
|
|
|
+
|
|
|
+on_message_publish(Msg) ->
|
|
|
+ {ok, Msg}.
|
|
|
+
|
|
|
+%% @private
|
|
|
+delayed_mid(undefined) ->
|
|
|
+ emqx_guid:gen();
|
|
|
+delayed_mid(MsgId) -> MsgId.
|
|
|
+
|
|
|
+%%--------------------------------------------------------------------
|
|
|
+%% 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).
|
|
|
+store(DelayedMsg) ->
|
|
|
+ gen_server:call(?SERVER, {store, DelayedMsg}, infinity).
|
|
|
+
|
|
|
+%%--------------------------------------------------------------------
|
|
|
+%% gen_server callback
|
|
|
+%%--------------------------------------------------------------------
|
|
|
+
|
|
|
+init([]) ->
|
|
|
+ ok = ekka_mnesia:create_table(?TAB, [
|
|
|
+ {type, ordered_set},
|
|
|
+ {disc_copies, [node()]},
|
|
|
+ {local_content, true},
|
|
|
+ {record_name, delayed_message},
|
|
|
+ {attributes, record_info(fields, delayed_message)}]),
|
|
|
+ ok = ekka_mnesia:copy_table(?TAB, disc_copies),
|
|
|
+ {ok, ensure_publish_timer(#{timer => undefined, publish_at => 0})}.
|
|
|
+
|
|
|
+handle_call({store, DelayedMsg = #delayed_message{key = Key}}, _From, State) ->
|
|
|
+ ok = mnesia:dirty_write(?TAB, DelayedMsg),
|
|
|
+ emqx_metrics:set('messages.delayed', delayed_count()),
|
|
|
+ {reply, ok, ensure_publish_timer(Key, State)};
|
|
|
+
|
|
|
+handle_call(Req, _From, State) ->
|
|
|
+ ?LOG(error, "[Delayed] Unexpected call: ~p", [Req]),
|
|
|
+ {reply, ignored, State}.
|
|
|
+
|
|
|
+handle_cast(Msg, State) ->
|
|
|
+ ?LOG(error, "[Delayed] Unexpected cast: ~p", [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) -> mnesia:dirty_delete(?TAB, Key) end, DeletedKeys),
|
|
|
+ emqx_metrics:set('messages.delayed', delayed_count()),
|
|
|
+ {noreply, ensure_publish_timer(State#{timer := undefined, publish_at := 0})};
|
|
|
+
|
|
|
+handle_info(Info, State) ->
|
|
|
+ ?LOG(error, "[Delayed] Unexpected info: ~p", [Info]),
|
|
|
+ {noreply, State}.
|
|
|
+
|
|
|
+terminate(_Reason, #{timer := TRef}) ->
|
|
|
+ emqx_misc:cancel_timer(TRef).
|
|
|
+
|
|
|
+code_change(_OldVsn, State, _Extra) ->
|
|
|
+ {ok, State}.
|
|
|
+
|
|
|
+%%--------------------------------------------------------------------
|
|
|
+%% Internal functions
|
|
|
+%%--------------------------------------------------------------------
|
|
|
+
|
|
|
+%% 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_broker: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).
|
|
|
+
|