fix pub/sub

This commit is contained in:
anlicheng 2026-05-11 14:20:36 +08:00
parent 748a4c2d9a
commit 4881308a69

View File

@ -2,7 +2,52 @@
%%% @author anlicheng
%%% @copyright (C) 2025, <COMPANY>
%%% @doc
%%% Endpoint
%%%
%%% publish
%%% exact wildcard
%%%
%%% <ul>
%%% <li> endpoint_subscription_exactpublish RouteKey
%%% ets:lookup/2</li>
%%% <li> ETS trietrie
%%% 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 idsubscribe/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_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);
(_) ->
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, Subscribers),
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)].