fix pub/sub

src/endpoint/endpoint_subscription.erl

@@ -2,7 +2,52 @@
 %%% @author anlicheng
 %%% @copyright (C) 2025, <COMPANY>
 %%% @doc
+%%% Endpoint 数据分发订阅索引。
 %%%
+%%% 当前实现面向较大订阅量场景，publish 热路径不再全量遍历订阅表，而是把
+%%% 订阅拆成 exact 和 wildcard 两类：
+%%%
+%%% <ul>
+%%%   <li>精确订阅写入 endpoint_subscription_exact，publish 时按 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 node，root 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
 %%%-------------------------------------------------------------------
@@ -14,13 +59,18 @@
 %% API
 -export([start_link/0]).
 -export([subscribe/2, unsubscribe/2, publish/2, get_subscribers/0]).
--export([match_components/2, is_valid_components/1, of_components/1]).
+-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(SUBSCRIBER_TAB, endpoint_subscription_subscribers).
+-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, {
@@ -36,7 +86,12 @@
 }).
 
 -record(state, {
-    tid :: ets:tid()
+    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()
 }).
 
 %%%===================================================================
@@ -51,23 +106,22 @@ subscribe(Topic, SubscriberPid) when is_binary(Topic), is_pid(SubscriberPid) ->
 unsubscribe(Topic, SubscriberPid) when is_binary(Topic), is_pid(SubscriberPid) ->
     gen_server:call(?SERVER, {unsubscribe, Topic, SubscriberPid}).
 
--spec get_subscribers() -> {ok, Subscribers :: map()}.
+-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(?SUBSCRIBER_TAB) of
+    case ets:info(?EXACT_TAB) of
         undefined ->
             ok;
         _ ->
-            Subscribers = ets:tab2list(?SUBSCRIBER_TAB),
+            MatchedSubscribers = match_route_key(RouteKey),
-            MatchedSubscribers = match_subscribers(Subscribers, RouteKey),
             lists:foreach(fun(#subscriber{subscriber_pid = SubscriberPid}) ->
                 endpoint:forward(SubscriberPid, Content)
             end, MatchedSubscribers),
             maybe_log_unmatched_publish(RouteKey, Content, MatchedSubscribers),
-            logger:debug("[efka_subscription] route_key: ~p, metric: ~p, match subscribers: ~p", [RouteKey, Content, MatchedSubscribers]),
+            logger:debug("[endpoint_subscription] route_key: ~p, match_count: ~p", [RouteKey, length(MatchedSubscribers)]),
             ok
     end.
 
@@ -88,8 +142,18 @@ start_link() ->
     {stop, Reason :: term()} | ignore).
 init([]) ->
     ok = iot_log:set_metadata(),
-    Tid = ets:new(?SUBSCRIBER_TAB, [named_table, protected, bag, {keypos, 2}]),
+    ExactTid = ets:new(?EXACT_TAB, [named_table, protected, bag, {read_concurrency, true}]),
-    {ok, #state{tid = Tid}}.
+    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
@@ -102,20 +166,19 @@ init([]) ->
     {stop, Reason :: term(), Reply :: term(), NewState :: #state{}} |
     {stop, Reason :: term(), NewState :: #state{}}).
 %%  同一个SubscriberPid只能订阅同一个topic一次
-handle_call(get_subscribers, _From, State = #state{tid = Tid}) ->
+handle_call(get_subscribers, _From, State = #state{exact_tid = ExactTid, trie_sub_tid = TrieSubTid}) ->
-    Subscribers = ets:tab2list(Tid),
+    Subscribers = exact_subscribers(ExactTid) ++ trie_subscribers(TrieSubTid),
     {reply, {ok, Subscribers}, State};
 
-handle_call({subscribe, Topic, SubscriberPid}, _From, State = #state{tid = Tid}) ->
+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(Tid, Topic, SubscriberPid) of
+            case has_subscription(ReverseTid, Topic, SubscriberPid) of
                 true ->
                     {reply, ok, State};
                 false ->
-                    %% 建立到SubscriberPid的monitor，进程退出需要清理订阅
+                    {MonitorRef, State1} = ensure_pid_monitor(SubscriberPid, State),
-                    MonitorRef = erlang:monitor(process, SubscriberPid),
                     Sub = #subscriber{
                         topic = Topic,
                         subscriber_pid = SubscriberPid,
@@ -123,23 +186,19 @@ handle_call({subscribe, Topic, SubscriberPid}, _From, State = #state{tid = Tid})
                         monitor_ref = MonitorRef,
                         order = order_num(Components)
                     },
-                    true = ets:insert(Tid, Sub),
+                    State2 = insert_subscription(Topic, Components, SubscriberPid, Sub, State1),
-                    {reply, ok, State}
+                    {reply, ok, State2}
             end;
         false ->
             {reply, {error, <<"invalid topic name">>}, State}
     end;
 
-handle_call({unsubscribe, Topic, SubscriberPid}, _From, State = #state{tid = Tid}) ->
+handle_call({unsubscribe, Topic, SubscriberPid}, _From, State = #state{reverse_tid = ReverseTid}) ->
-    Removed = [Sub || Sub = #subscriber{subscriber_pid = SubscriberPid0} <- ets:lookup(Tid, Topic),
+    Removed = [Reverse || Reverse = {SubscriberPid0, Topic0, _Kind, _NodeId, _Sub} <- ets:lookup(ReverseTid, SubscriberPid),
-        SubscriberPid =:= SubscriberPid0],
+        SubscriberPid =:= SubscriberPid0, Topic =:= Topic0],
-    lists:foreach(fun(#subscriber{monitor_ref = MonitorRef}) when is_reference(MonitorRef) ->
+    lists:foreach(fun(Reverse) -> delete_subscription(Reverse, State) end, Removed),
-        erlang:demonitor(MonitorRef, [flush]);
+    State1 = release_pid_monitor(SubscriberPid, length(Removed), State),
-                     (_) ->
+    {reply, ok, State1}.
-        ok
-    end, Removed),
-    lists:foreach(fun(Sub) -> true = ets:delete_object(Tid, Sub) end, Removed),
-    {reply, ok, State}.
 
 %% @private
 %% @doc Handling cast messages
@@ -156,18 +215,20 @@ handle_cast(_Request, State = #state{}) ->
     {noreply, NewState :: #state{}} |
     {noreply, NewState :: #state{}, timeout() | hibernate} |
     {stop, Reason :: term(), NewState :: #state{}}).
-handle_info({'DOWN', _Ref, process, SubscriberPid, Reason}, State = #state{tid = Tid}) ->
+handle_info({'DOWN', MonitorRef, process, SubscriberPid, Reason}, State = #state{reverse_tid = ReverseTid, pid_tid = PidTid}) ->
-    logger:debug("[efka_subscription] subscriber: ~p, down with reason: ~p", [SubscriberPid, Reason]),
+    logger:debug("[endpoint_subscription] subscriber: ~p, down with reason: ~p", [SubscriberPid, Reason]),
-    Subscribers = ets:tab2list(Tid),
+    Removed = ets:lookup(ReverseTid, SubscriberPid),
-    lists:foreach(fun(Sub = #subscriber{subscriber_pid = Pid0}) when Pid0 =:= SubscriberPid ->
+    lists:foreach(fun(Reverse) -> delete_subscription(Reverse, State) end, Removed),
-        true = ets:delete_object(Tid, Sub);
+    case ets:lookup(PidTid, SubscriberPid) of
-                     (_) ->
+        [{SubscriberPid, MonitorRef, _Count}] ->
-        ok
+            ets:delete(PidTid, SubscriberPid);
-    end, Subscribers),
+        _ ->
+            ok
+    end,
     {noreply, State};
 
 handle_info(Info, State = #state{}) ->
-    logger:debug("[efka_subscription] get unknown info: ~p", [Info]),
+    logger:debug("[endpoint_subscription] get unknown info: ~p", [Info]),
     {noreply, State}.
 
 %% @private
@@ -192,14 +253,11 @@ code_change(_OldVsn, State = #state{}, _Extra) ->
 %%% Internal functions
 %%%===================================================================
 
-%% 查找满足条件订阅者
+-spec match_route_key(binary()) -> [#subscriber{}].
--spec match_subscribers(Subscribers :: [#subscriber{}], Topic :: binary()) -> [#subscriber{}].
+match_route_key(RouteKey) when is_binary(RouteKey) ->
-match_subscribers(Subscribers, Topic) when is_list(Subscribers), is_binary(Topic) ->
+    ExactSubs = [Sub || {_Topic, Sub} <- ets:lookup(?EXACT_TAB, RouteKey)],
-    Components = of_components(Topic),
+    TrieSubs = match_trie(of_components(RouteKey)),
-    Matched = lists:filter(fun(#subscriber{components = Components0}) ->
+    Sorted = lists:sort(fun compare_subscriber/2, ExactSubs ++ TrieSubs),
-        match_components(Components0, Components)
-    end, Subscribers),
-    Sorted = lists:sort(fun compare_subscriber/2, Matched),
     dedupe_subscribers(Sorted).
 
 -spec maybe_log_unmatched_publish(binary(), binary(), [#subscriber{}]) -> ok.
@@ -208,22 +266,6 @@ maybe_log_unmatched_publish(RouteKey, Content, []) ->
 maybe_log_unmatched_publish(_RouteKey, _Content, _MatchedSubscribers) ->
     ok.
 
-%% 开始对比订阅的topic和发布的topic的Components信息
-%% *表示单级匹配，+表示多级匹配；+只能出现一次，并且只能在末尾
--spec match_components(list(), list()) -> boolean().
-match_components(A, B) when is_list(A), is_list(B) ->
-    match_components(A, B, false).
-match_components([<<"+">>], [_|_], _) ->
-    true;
-match_components([], [], _) ->
-    true;
-match_components([<<"*">>|T0], [_|T1], _) ->
-    match_components(T0, T1, false);
-match_components([C0|T0], [C0|T1], _) ->
-    match_components(T0, T1, false);
-match_components(_, _, _) ->
-    false.
-
 -spec of_components(Topic :: binary()) -> [binary()].
 of_components(Topic) when is_binary(Topic) ->
     binary:split(Topic, <<$/>>, [global]).
@@ -248,11 +290,11 @@ order_num([_|Tail]) ->
     order_num(Tail).
 
 -spec has_subscription(ets:tid(), binary(), pid()) -> boolean().
-has_subscription(Tid, Topic, SubscriberPid) ->
+has_subscription(ReverseTid, Topic, SubscriberPid) ->
-    Subscribers = ets:lookup(Tid, Topic),
+    Subscriptions = ets:lookup(ReverseTid, SubscriberPid),
-    lists:any(fun(#subscriber{topic = Topic0, subscriber_pid = SubscriberPid0}) ->
+    lists:any(fun({SubscriberPid0, Topic0, _Kind, _NodeId, _Sub}) ->
         Topic =:= Topic0 andalso SubscriberPid =:= SubscriberPid0
-    end, Subscribers).
+    end, Subscriptions).
 
 -spec compare_subscriber(#subscriber{}, #subscriber{}) -> boolean().
 compare_subscriber(#subscriber{order = Order0, topic = Topic0}, #subscriber{order = Order1, topic = Topic1}) ->
@@ -274,3 +316,142 @@ dedupe_subscribers(Subscribers) ->
         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)].