%%%------------------------------------------------------------------- %%% @author anlicheng %%% @copyright (C) 2025, %%% @doc %%% Endpoint 数据分发订阅索引。 %%% %%% 当前实现面向较大订阅量场景，publish 热路径不再全量遍历订阅表，而是把 %%% 订阅拆成 exact 和 wildcard 两类： %%% %%% %%% %%% 所有 matcher 共用同一个 root node，root node id 固定为 0。新 trie 节点使用 %%% 单调递增的 integer node id。subscribe/unsubscribe 通过本 gen_server 串行写 ETS； %%% publish/2 直接并发读取 ETS，避免中心 gen_server 成为数据分发瓶颈。 %%% %%% 匹配规则： %%% %%% %%% %%% publish/2 的匹配流程： %%% %%%
    %%%
  1. 从 exact 表直接查找 RouteKey 对应的订阅。
    2. %%%
  2. 把 RouteKey 按 "/" 拆分后，从 trie root 开始逐层匹配。每层最多查找 %%% 当前 segment 分支和 <<"*">> 分支。
    3. %%%
  3. 每消费一层前，收集当前候选节点上的 plus 订阅；全部 segment 消费完后， %%% 收集候选节点上的 exact 订阅。
    4. %%%
  4. 合并 exact 和 trie 命中结果，按原有优先级排序，再按 SubscriberPid 去重。 %%% 同一个 pid 同时被多个 matcher 命中时，保留优先级最高的订阅。
    5. %%%
  5. 对最终订阅者执行 endpoint:forward/2。
    6. %%%
%%% @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)].