From 4881308a691f49a9e391ec6cd0339d66335ee35e Mon Sep 17 00:00:00 2001 From: anlicheng <244108715@qq.com> Date: Mon, 11 May 2026 14:20:36 +0800 Subject: [PATCH] fix pub/sub --- src/endpoint/endpoint_subscription.erl | 311 +++++++++++++++++++------ 1 file changed, 246 insertions(+), 65 deletions(-) diff --git a/src/endpoint/endpoint_subscription.erl b/src/endpoint/endpoint_subscription.erl index 18b9754..6197962 100644 --- a/src/endpoint/endpoint_subscription.erl +++ b/src/endpoint/endpoint_subscription.erl @@ -2,7 +2,52 @@ %%% @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 %%%------------------------------------------------------------------- @@ -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_subscribers(Subscribers, RouteKey), + 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("[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}]), - {ok, #state{tid = Tid}}. + 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 @@ -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}) -> - Subscribers = ets:tab2list(Tid), +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{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 = erlang:monitor(process, SubscriberPid), + {MonitorRef, State1} = ensure_pid_monitor(SubscriberPid, State), 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), - {reply, ok, State} + 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{tid = Tid}) -> - Removed = [Sub || Sub = #subscriber{subscriber_pid = SubscriberPid0} <- ets:lookup(Tid, Topic), - SubscriberPid =:= SubscriberPid0], - lists:foreach(fun(#subscriber{monitor_ref = MonitorRef}) when is_reference(MonitorRef) -> - erlang:demonitor(MonitorRef, [flush]); - (_) -> - ok - end, Removed), - lists:foreach(fun(Sub) -> true = ets:delete_object(Tid, Sub) end, Removed), - {reply, ok, State}. +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 @@ -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}) -> - logger:debug("[efka_subscription] subscriber: ~p, down with reason: ~p", [SubscriberPid, Reason]), - Subscribers = ets:tab2list(Tid), - lists:foreach(fun(Sub = #subscriber{subscriber_pid = Pid0}) when Pid0 =:= SubscriberPid -> - true = ets:delete_object(Tid, Sub); - (_) -> - ok - end, Subscribers), +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("[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_subscribers(Subscribers :: [#subscriber{}], Topic :: binary()) -> [#subscriber{}]. -match_subscribers(Subscribers, Topic) when is_list(Subscribers), is_binary(Topic) -> - Components = of_components(Topic), - Matched = lists:filter(fun(#subscriber{components = Components0}) -> - match_components(Components0, Components) - end, Subscribers), - Sorted = lists:sort(fun compare_subscriber/2, Matched), +-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. @@ -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) -> - Subscribers = ets:lookup(Tid, Topic), - lists:any(fun(#subscriber{topic = Topic0, subscriber_pid = SubscriberPid0}) -> +has_subscription(ReverseTid, Topic, SubscriberPid) -> + Subscriptions = ets:lookup(ReverseTid, SubscriberPid), + 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)].