iot_cloud/src/endpoint/endpoint_subscription.erl
2026-05-11 14:20:36 +08:00

458 lines
19 KiB
Erlang
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

%%%-------------------------------------------------------------------
%%% @author anlicheng
%%% @copyright (C) 2025, <COMPANY>
%%% @doc
%%% Endpoint 数据分发订阅索引。
%%%
%%% 当前实现面向较大订阅量场景publish 热路径不再全量遍历订阅表,而是把
%%% 订阅拆成 exact 和 wildcard 两类:
%%%
%%% <ul>
%%% <li>精确订阅写入 endpoint_subscription_exactpublish 时按 RouteKey
%%% 直接 ets:lookup/2。</li>
%%% <li>包含通配符的订阅写入 ETS trie。trie 边存放在
%%% endpoint_subscription_trie_edge节点订阅存放在
%%% endpoint_subscription_trie_sub。</li>
%%% <li>endpoint_subscription_reverse 保存 SubscriberPid 到订阅对象的反向索引,
%%% 用于 unsubscribe 和进程 DOWN 清理。</li>
%%% <li>endpoint_subscription_pid 记录每个 SubscriberPid 的 monitor 和订阅计数,
%%% 同一个 pid 只 monitor 一次。</li>
%%% </ul>
%%%
%%% 所有 matcher 共用同一个 root noderoot node id 固定为 0。新 trie 节点使用
%%% 单调递增的 integer node id。subscribe/unsubscribe 通过本 gen_server 串行写 ETS
%%% publish/2 直接并发读取 ETS避免中心 gen_server 成为数据分发瓶颈。
%%%
%%% 匹配规则:
%%%
%%% <ul>
%%% <li>普通 segment 必须完全相等,例如 <<"device/a/temp">> 只精确匹配
%%% 同名订阅。</li>
%%% <li><<"*">> 是单级通配,只匹配一个 segment。例如
%%% <<"device/*/temp">> 匹配 <<"device/a/temp">>,不匹配
%%% <<"device/a/b/temp">>。</li>
%%% <li><<"+">> 是末尾多级通配,只允许出现在 matcher 最后一段,并且至少
%%% 匹配一个剩余 segment。例如 <<"device/+">> 匹配 <<"device/a">>
%%% 和 <<"device/a/temp">>,不匹配 <<"device">>。</li>
%%% </ul>
%%%
%%% publish/2 的匹配流程:
%%%
%%% <ol>
%%% <li>从 exact 表直接查找 RouteKey 对应的订阅。</li>
%%% <li>把 RouteKey 按 "/" 拆分后,从 trie root 开始逐层匹配。每层最多查找
%%% 当前 segment 分支和 <<"*">> 分支。</li>
%%% <li>每消费一层前,收集当前候选节点上的 plus 订阅;全部 segment 消费完后,
%%% 收集候选节点上的 exact 订阅。</li>
%%% <li>合并 exact 和 trie 命中结果,按原有优先级排序,再按 SubscriberPid 去重。
%%% 同一个 pid 同时被多个 matcher 命中时,保留优先级最高的订阅。</li>
%%% <li>对最终订阅者执行 endpoint:forward/2。</li>
%%% </ol>
%%% @end
%%% Created : 07. 11月 2025 16:27
%%%-------------------------------------------------------------------
-module(endpoint_subscription).
-author("anlicheng").
-behaviour(gen_server).
%% API
-export([start_link/0]).
-export([subscribe/2, unsubscribe/2, publish/2, get_subscribers/0]).
-export([is_valid_components/1, of_components/1]).
%% gen_server callbacks
-export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]).
-define(SERVER, ?MODULE).
-define(EXACT_TAB, endpoint_subscription_exact).
-define(TRIE_EDGE_TAB, endpoint_subscription_trie_edge).
-define(TRIE_SUB_TAB, endpoint_subscription_trie_sub).
-define(REVERSE_TAB, endpoint_subscription_reverse).
-define(PID_TAB, endpoint_subscription_pid).
-define(ROOT_NODE, 0).
%% 定义订阅者
-record(subscriber, {
topic :: binary(),
subscriber_pid :: pid(),
components = [],
monitor_ref :: undefined | reference(),
%% 优先级
%% 1. 完全匹配的topic优先级别最高
%% 2. 带 * 的订阅
%% 3. 带 + 的订阅
order :: integer()
}).
-record(state, {
exact_tid :: ets:tid(),
edge_tid :: ets:tid(),
trie_sub_tid :: ets:tid(),
reverse_tid :: ets:tid(),
pid_tid :: ets:tid(),
next_node_id = 1 :: pos_integer()
}).
%%%===================================================================
%%% API
%%%===================================================================
-spec subscribe(Topic :: binary(), SubscriberPid :: pid()) -> ok | {error, Reason :: binary()}.
subscribe(Topic, SubscriberPid) when is_binary(Topic), is_pid(SubscriberPid) ->
gen_server:call(?SERVER, {subscribe, Topic, SubscriberPid}).
-spec unsubscribe(Topic :: binary(), SubscriberPid :: pid()) -> ok.
unsubscribe(Topic, SubscriberPid) when is_binary(Topic), is_pid(SubscriberPid) ->
gen_server:call(?SERVER, {unsubscribe, Topic, SubscriberPid}).
-spec get_subscribers() -> {ok, Subscribers :: list()}.
get_subscribers() ->
gen_server:call(?SERVER, get_subscribers).
-spec publish(RouteKey :: binary(), Content :: binary()) -> ok.
publish(RouteKey, Content) when is_binary(RouteKey), is_binary(Content) ->
case ets:info(?EXACT_TAB) of
undefined ->
ok;
_ ->
MatchedSubscribers = match_route_key(RouteKey),
lists:foreach(fun(#subscriber{subscriber_pid = SubscriberPid}) ->
endpoint:forward(SubscriberPid, Content)
end, MatchedSubscribers),
maybe_log_unmatched_publish(RouteKey, Content, MatchedSubscribers),
logger:debug("[endpoint_subscription] route_key: ~p, match_count: ~p", [RouteKey, length(MatchedSubscribers)]),
ok
end.
%% @doc Spawns the server and registers the local name (unique)
-spec(start_link() ->
{ok, Pid :: pid()} | ignore | {error, Reason :: term()}).
start_link() ->
gen_server:start_link({local, ?SERVER}, ?MODULE, [], []).
%%%===================================================================
%%% gen_server callbacks
%%%===================================================================
%% @private
%% @doc Initializes the server
-spec(init(Args :: term()) ->
{ok, State :: #state{}} | {ok, State :: #state{}, timeout() | hibernate} |
{stop, Reason :: term()} | ignore).
init([]) ->
ok = iot_log:set_metadata(),
ExactTid = ets:new(?EXACT_TAB, [named_table, protected, bag, {read_concurrency, true}]),
EdgeTid = ets:new(?TRIE_EDGE_TAB, [named_table, protected, set, {read_concurrency, true}]),
TrieSubTid = ets:new(?TRIE_SUB_TAB, [named_table, protected, bag, {read_concurrency, true}]),
ReverseTid = ets:new(?REVERSE_TAB, [named_table, protected, bag]),
PidTid = ets:new(?PID_TAB, [named_table, protected, set]),
{ok, #state{
exact_tid = ExactTid,
edge_tid = EdgeTid,
trie_sub_tid = TrieSubTid,
reverse_tid = ReverseTid,
pid_tid = PidTid
}}.
%% @private
%% @doc Handling call messages
-spec(handle_call(Request :: term(), From :: {pid(), Tag :: term()},
State :: #state{}) ->
{reply, Reply :: term(), NewState :: #state{}} |
{reply, Reply :: term(), NewState :: #state{}, timeout() | hibernate} |
{noreply, NewState :: #state{}} |
{noreply, NewState :: #state{}, timeout() | hibernate} |
{stop, Reason :: term(), Reply :: term(), NewState :: #state{}} |
{stop, Reason :: term(), NewState :: #state{}}).
%% 同一个SubscriberPid只能订阅同一个topic一次
handle_call(get_subscribers, _From, State = #state{exact_tid = ExactTid, trie_sub_tid = TrieSubTid}) ->
Subscribers = exact_subscribers(ExactTid) ++ trie_subscribers(TrieSubTid),
{reply, {ok, Subscribers}, State};
handle_call({subscribe, Topic, SubscriberPid}, _From, State = #state{reverse_tid = ReverseTid}) ->
Components = of_components(Topic),
case is_valid_components(Components) of
true ->
case has_subscription(ReverseTid, Topic, SubscriberPid) of
true ->
{reply, ok, State};
false ->
{MonitorRef, State1} = ensure_pid_monitor(SubscriberPid, State),
Sub = #subscriber{
topic = Topic,
subscriber_pid = SubscriberPid,
components = Components,
monitor_ref = MonitorRef,
order = order_num(Components)
},
State2 = insert_subscription(Topic, Components, SubscriberPid, Sub, State1),
{reply, ok, State2}
end;
false ->
{reply, {error, <<"invalid topic name">>}, State}
end;
handle_call({unsubscribe, Topic, SubscriberPid}, _From, State = #state{reverse_tid = ReverseTid}) ->
Removed = [Reverse || Reverse = {SubscriberPid0, Topic0, _Kind, _NodeId, _Sub} <- ets:lookup(ReverseTid, SubscriberPid),
SubscriberPid =:= SubscriberPid0, Topic =:= Topic0],
lists:foreach(fun(Reverse) -> delete_subscription(Reverse, State) end, Removed),
State1 = release_pid_monitor(SubscriberPid, length(Removed), State),
{reply, ok, State1}.
%% @private
%% @doc Handling cast messages
-spec(handle_cast(Request :: term(), State :: #state{}) ->
{noreply, NewState :: #state{}} |
{noreply, NewState :: #state{}, timeout() | hibernate} |
{stop, Reason :: term(), NewState :: #state{}}).
handle_cast(_Request, State = #state{}) ->
{noreply, State}.
%% @private
%% @doc Handling all non call/cast messages
-spec(handle_info(Info :: timeout() | term(), State :: #state{}) ->
{noreply, NewState :: #state{}} |
{noreply, NewState :: #state{}, timeout() | hibernate} |
{stop, Reason :: term(), NewState :: #state{}}).
handle_info({'DOWN', MonitorRef, process, SubscriberPid, Reason}, State = #state{reverse_tid = ReverseTid, pid_tid = PidTid}) ->
logger:debug("[endpoint_subscription] subscriber: ~p, down with reason: ~p", [SubscriberPid, Reason]),
Removed = ets:lookup(ReverseTid, SubscriberPid),
lists:foreach(fun(Reverse) -> delete_subscription(Reverse, State) end, Removed),
case ets:lookup(PidTid, SubscriberPid) of
[{SubscriberPid, MonitorRef, _Count}] ->
ets:delete(PidTid, SubscriberPid);
_ ->
ok
end,
{noreply, State};
handle_info(Info, State = #state{}) ->
logger:debug("[endpoint_subscription] get unknown info: ~p", [Info]),
{noreply, State}.
%% @private
%% @doc This function is called by a gen_server when it is about to
%% terminate. It should be the opposite of Module:init/1 and do any
%% necessary cleaning up. When it returns, the gen_server terminates
%% with Reason. The return value is ignored.
-spec(terminate(Reason :: (normal | shutdown | {shutdown, term()} | term()),
State :: #state{}) -> term()).
terminate(_Reason, _State = #state{}) ->
ok.
%% @private
%% @doc Convert process state when code is changed
-spec(code_change(OldVsn :: term() | {down, term()}, State :: #state{},
Extra :: term()) ->
{ok, NewState :: #state{}} | {error, Reason :: term()}).
code_change(_OldVsn, State = #state{}, _Extra) ->
{ok, State}.
%%%===================================================================
%%% Internal functions
%%%===================================================================
-spec match_route_key(binary()) -> [#subscriber{}].
match_route_key(RouteKey) when is_binary(RouteKey) ->
ExactSubs = [Sub || {_Topic, Sub} <- ets:lookup(?EXACT_TAB, RouteKey)],
TrieSubs = match_trie(of_components(RouteKey)),
Sorted = lists:sort(fun compare_subscriber/2, ExactSubs ++ TrieSubs),
dedupe_subscribers(Sorted).
-spec maybe_log_unmatched_publish(binary(), binary(), [#subscriber{}]) -> ok.
maybe_log_unmatched_publish(RouteKey, Content, []) ->
endpoint_log:unmatched_publish(RouteKey, Content);
maybe_log_unmatched_publish(_RouteKey, _Content, _MatchedSubscribers) ->
ok.
-spec of_components(Topic :: binary()) -> [binary()].
of_components(Topic) when is_binary(Topic) ->
binary:split(Topic, <<$/>>, [global]).
is_valid_components([]) ->
true;
is_valid_components([<<$+>>|T]) ->
length(T) =:= 0;
is_valid_components([<<$*>>|T]) ->
is_valid_components(T);
is_valid_components([_|T]) ->
is_valid_components(T).
-spec order_num(Components :: list()) -> integer().
order_num([]) ->
1;
order_num([<<$*>>|_]) ->
2;
order_num([<<$+>>|_]) ->
3;
order_num([_|Tail]) ->
order_num(Tail).
-spec has_subscription(ets:tid(), binary(), pid()) -> boolean().
has_subscription(ReverseTid, Topic, SubscriberPid) ->
Subscriptions = ets:lookup(ReverseTid, SubscriberPid),
lists:any(fun({SubscriberPid0, Topic0, _Kind, _NodeId, _Sub}) ->
Topic =:= Topic0 andalso SubscriberPid =:= SubscriberPid0
end, Subscriptions).
-spec compare_subscriber(#subscriber{}, #subscriber{}) -> boolean().
compare_subscriber(#subscriber{order = Order0, topic = Topic0}, #subscriber{order = Order1, topic = Topic1}) ->
case Order0 =:= Order1 of
true ->
Topic0 =< Topic1;
false ->
Order0 < Order1
end.
-spec dedupe_subscribers([#subscriber{}]) -> [#subscriber{}].
dedupe_subscribers(Subscribers) ->
{_, Result} = lists:foldl(fun(S = #subscriber{subscriber_pid = SubscriberPid}, {Seen, Acc}) ->
case sets:is_element(SubscriberPid, Seen) of
true ->
{Seen, Acc};
false ->
{sets:add_element(SubscriberPid, Seen), [S | Acc]}
end
end, {sets:new(), []}, Subscribers),
lists:reverse(Result).
-spec insert_subscription(binary(), [binary()], pid(), #subscriber{}, #state{}) -> #state{}.
insert_subscription(Topic, Components, SubscriberPid, Sub, State = #state{exact_tid = ExactTid, trie_sub_tid = TrieSubTid, reverse_tid = ReverseTid}) ->
case has_wildcard(Components) of
false ->
true = ets:insert(ExactTid, {Topic, Sub}),
true = ets:insert(ReverseTid, {SubscriberPid, Topic, exact, undefined, Sub}),
State;
true ->
{Kind, NodeId, State1} = insert_trie_subscription(Components, State),
true = ets:insert(TrieSubTid, {{NodeId, subscription_match_type(Kind)}, Sub}),
true = ets:insert(ReverseTid, {SubscriberPid, Topic, Kind, NodeId, Sub}),
State1
end.
-spec insert_trie_subscription([binary()], #state{}) -> {trie_exact | trie_plus, non_neg_integer(), #state{}}.
insert_trie_subscription(Components, State) ->
case Components of
[] ->
{trie_exact, ?ROOT_NODE, State};
_ ->
case lists:last(Components) of
<<"+">> ->
Prefix = lists:sublist(Components, length(Components) - 1),
{NodeId, State1} = ensure_path(Prefix, State),
{trie_plus, NodeId, State1};
_ ->
{NodeId, State1} = ensure_path(Components, State),
{trie_exact, NodeId, State1}
end
end.
-spec subscription_match_type(trie_exact | trie_plus) -> exact | plus.
subscription_match_type(trie_exact) ->
exact;
subscription_match_type(trie_plus) ->
plus.
-spec ensure_path([binary()], #state{}) -> {non_neg_integer(), #state{}}.
ensure_path(Components, State) ->
ensure_path(?ROOT_NODE, Components, State).
-spec ensure_path(non_neg_integer(), [binary()], #state{}) -> {non_neg_integer(), #state{}}.
ensure_path(NodeId, [], State) ->
{NodeId, State};
ensure_path(NodeId, [Segment | Rest], State = #state{edge_tid = EdgeTid, next_node_id = NextNodeId}) ->
EdgeKey = {NodeId, Segment},
case ets:lookup(EdgeTid, EdgeKey) of
[{EdgeKey, ChildNodeId}] ->
ensure_path(ChildNodeId, Rest, State);
[] ->
true = ets:insert(EdgeTid, {EdgeKey, NextNodeId}),
ensure_path(NextNodeId, Rest, State#state{next_node_id = NextNodeId + 1})
end.
-spec delete_subscription(tuple(), #state{}) -> ok.
delete_subscription(Reverse = {_SubscriberPid, Topic, exact, undefined, Sub}, #state{exact_tid = ExactTid, reverse_tid = ReverseTid}) ->
true = ets:delete_object(ExactTid, {Topic, Sub}),
true = ets:delete_object(ReverseTid, Reverse),
ok;
delete_subscription(Reverse = {_SubscriberPid, _Topic, Kind, NodeId, Sub}, #state{trie_sub_tid = TrieSubTid, reverse_tid = ReverseTid}) ->
true = ets:delete_object(TrieSubTid, {{NodeId, subscription_match_type(Kind)}, Sub}),
true = ets:delete_object(ReverseTid, Reverse),
ok.
-spec ensure_pid_monitor(pid(), #state{}) -> {reference(), #state{}}.
ensure_pid_monitor(SubscriberPid, State = #state{pid_tid = PidTid}) ->
case ets:lookup(PidTid, SubscriberPid) of
[{SubscriberPid, MonitorRef, Count}] ->
true = ets:insert(PidTid, {SubscriberPid, MonitorRef, Count + 1}),
{MonitorRef, State};
[] ->
MonitorRef = erlang:monitor(process, SubscriberPid),
true = ets:insert(PidTid, {SubscriberPid, MonitorRef, 1}),
{MonitorRef, State}
end.
-spec release_pid_monitor(pid(), non_neg_integer(), #state{}) -> #state{}.
release_pid_monitor(_SubscriberPid, 0, State) ->
State;
release_pid_monitor(SubscriberPid, RemovedCount, State = #state{pid_tid = PidTid}) ->
case ets:lookup(PidTid, SubscriberPid) of
[{SubscriberPid, MonitorRef, Count}] when Count =< RemovedCount ->
erlang:demonitor(MonitorRef, [flush]),
ets:delete(PidTid, SubscriberPid),
State;
[{SubscriberPid, MonitorRef, Count}] ->
true = ets:insert(PidTid, {SubscriberPid, MonitorRef, Count - RemovedCount}),
State;
[] ->
State
end.
-spec match_trie([binary()]) -> [#subscriber{}].
match_trie(Components) ->
match_trie([?ROOT_NODE], Components, []).
-spec match_trie([non_neg_integer()], [binary()], [#subscriber{}]) -> [#subscriber{}].
match_trie(NodeIds, [], Acc) ->
collect_trie_subscribers(NodeIds, exact) ++ Acc;
match_trie(NodeIds, [Segment | Rest], Acc) ->
PlusSubs = collect_trie_subscribers(NodeIds, plus),
NextNodeIds = next_trie_nodes(NodeIds, Segment),
match_trie(NextNodeIds, Rest, PlusSubs ++ Acc).
-spec collect_trie_subscribers([non_neg_integer()], exact | plus) -> [#subscriber{}].
collect_trie_subscribers(NodeIds, MatchType) ->
lists:flatmap(fun(NodeId) ->
[Sub || {{NodeId0, MatchType0}, Sub} <- ets:lookup(?TRIE_SUB_TAB, {NodeId, MatchType}),
NodeId0 =:= NodeId, MatchType0 =:= MatchType]
end, NodeIds).
-spec next_trie_nodes([non_neg_integer()], binary()) -> [non_neg_integer()].
next_trie_nodes(NodeIds, Segment) ->
lists:usort(lists:flatmap(fun(NodeId) ->
lookup_child(NodeId, Segment) ++ lookup_child(NodeId, <<"*">>)
end, NodeIds)).
-spec lookup_child(non_neg_integer(), binary()) -> [non_neg_integer()].
lookup_child(NodeId, Segment) ->
EdgeKey = {NodeId, Segment},
case ets:lookup(?TRIE_EDGE_TAB, EdgeKey) of
[{EdgeKey, ChildNodeId}] ->
[ChildNodeId];
[] ->
[]
end.
-spec has_wildcard([binary()]) -> boolean().
has_wildcard(Components) ->
lists:any(fun(Component) -> Component =:= <<"*">> orelse Component =:= <<"+">> end, Components).
-spec exact_subscribers(ets:tid()) -> [#subscriber{}].
exact_subscribers(ExactTid) ->
[Sub || {_Topic, Sub} <- ets:tab2list(ExactTid)].
-spec trie_subscribers(ets:tid()) -> [#subscriber{}].
trie_subscribers(TrieSubTid) ->
[Sub || {_Key, Sub} <- ets:tab2list(TrieSubTid)].