diff --git a/src/endpoint/endpoint.erl b/src/endpoint/endpoint.erl index c1984e9..12070ee 100644 --- a/src/endpoint/endpoint.erl +++ b/src/endpoint/endpoint.erl @@ -52,11 +52,11 @@ forward(Pid, Metric) when is_pid(Pid), is_binary(Metric) -> gen_server:cast(Pid, {forward, Metric}). reload(Pid, NEndpoint = #endpoint{}) when is_pid(Pid) -> - gen_statem:cast(Pid, {reload, NEndpoint}). + gen_server:cast(Pid, {reload, NEndpoint}). -spec clean_up(Pid :: pid()) -> ok. clean_up(Pid) when is_pid(Pid) -> - gen_server:call(Pid, clean_up, 5000). + gen_server:cast(Pid, cleanup). -spec get_protocol(Endpoint :: #endpoint{}) -> atom(). get_protocol(#endpoint{config = #http_endpoint{}}) -> @@ -118,11 +118,24 @@ parse_config(<<"mqtt">>, #{<<"host">> := Host, <<"port">> := Port0, <<"client_id end; parse_config(<<"http">>, C = #{<<"url">> := Url, <<"pool_size">> := PoolSize}) -> Token = maps:get(<<"token">>, C, <<>>), - {ok, #http_endpoint{ - url = Url, - token = Token, - pool_size = PoolSize - }}; + Errors = lists:filtermap(fun(Term) -> + case check_http_argument(Term) of + ok -> + false; + {error, Error} -> + {true, Error} + end + end, [{url, Url}, {token, Token}, {pool_size, PoolSize}]), + case Errors =:= [] of + true -> + {ok, #http_endpoint{ + url = Url, + token = Token, + pool_size = PoolSize + }}; + false -> + {error, Errors} + end; parse_config(<<"kafka">>, #{<<"sasl_config">> := #{<<"username">> := Username, <<"password">> := Password, <<"mechanism">> := Mechanism0}, <<"bootstrap_servers">> := BootstrapServers, <<"topic">> := Topic}) -> Errors = lists:filtermap(fun(Term) -> case check_kafka_argument(Term) of @@ -243,6 +256,29 @@ check_kafka_argument({bootstrap_servers, BootstrapServers}) -> {error, <<"bootstrap_servers is empty">>} end. +-spec check_http_argument(tuple()) -> ok | {error, Reason :: binary()}. +check_http_argument({url, Url}) -> + case is_binary(Url) andalso Url /= <<>> of + true -> + ok; + false -> + {error, <<"url is empty">>} + end; +check_http_argument({token, Token}) -> + case is_binary(Token) of + true -> + ok; + false -> + {error, <<"token invalid">>} + end; +check_http_argument({pool_size, PoolSize}) -> + case is_integer(PoolSize) andalso PoolSize > 0 of + true -> + ok; + false -> + {error, <<"pool_size invalid">>} + end. + -spec check_mqtt_argument(tuple()) -> ok | {error, Reason :: binary()}. check_mqtt_argument({host, Host}) -> case Host /= <<>> of @@ -285,4 +321,4 @@ check_mqtt_argument({qos, Qos}) -> ok; false -> {error, <<"qos invalid">>} - end. \ No newline at end of file + end. diff --git a/src/endpoint/endpoint_buffer.erl b/src/endpoint/endpoint_buffer.erl index 4368a53..ac5cc02 100644 --- a/src/endpoint/endpoint_buffer.erl +++ b/src/endpoint/endpoint_buffer.erl @@ -14,13 +14,13 @@ %% 消息重发间隔 -define(RETRY_INTERVAL, 5000). --export([new/2, append/2, trigger_next/1, trigger_n/1, cleanup/1, ack/2, stat/1]). +-export([new/2, append/2, trigger_next/1, trigger_n/1, cleanup/1, ack/2, stat/1, resize/2]). -export_type([buffer/0]). -record(buffer, { endpoint :: #endpoint{}, next_id = 1 :: integer(), - %% 当前数据所在的游标 + %% 最近一次派发的数据游标 cursor = 0 :: integer(), %% ets存储的引用 tid :: ets:tid(), @@ -36,7 +36,8 @@ -record(north_data, { id :: integer(), - tuple :: any() + tuple :: any(), + inflight = false :: boolean() }). -type buffer() :: #buffer{}. @@ -51,7 +52,7 @@ new(Endpoint = #endpoint{id = Id}, WindowSize) when is_integer(WindowSize), Wind #buffer{cursor = 0, tid = Tid, timer_pid = TimerPid, endpoint = Endpoint, window_size = WindowSize}. --spec append(tuple(), Buffer :: #buffer{}) -> NBuffer :: #buffer{}. +-spec append(Tuple :: any(), Buffer :: #buffer{}) -> NBuffer :: #buffer{}. append(Tuple, Buffer = #buffer{tid = Tid, next_id = NextId, window_size = WindowSize, flight_num = FlightNum}) -> NorthData = #north_data{id = NextId, tuple = Tuple}, true = ets:insert(Tid, NorthData), @@ -70,29 +71,43 @@ trigger_n(Buffer = #buffer{window_size = WindowSize}) -> %% 触发读取下一条数据 -spec trigger_next(Buffer :: #buffer{}) -> NBuffer :: #buffer{}. -trigger_next(Buffer = #buffer{tid = Tid, cursor = Cursor, timer_pid = TimerPid, flight_num = FlightNum}) -> - case ets:first(Tid) of - '$end_of_table' -> +trigger_next(Buffer = #buffer{tid = Tid, cursor = Cursor, timer_pid = TimerPid, flight_num = FlightNum, window_size = WindowSize}) -> + case FlightNum < WindowSize of + false -> Buffer; - Key -> - [#north_data{id = Id, tuple = Tuple}|_] = ets:take(Tid, Key), - ReceiverPid = self(), - ReceiverPid ! {next_data, Id, Tuple}, - endpoint_timer:task(TimerPid, Id, fun() -> ReceiverPid ! {next_data, Id, Tuple} end), - Buffer#buffer{flight_num = FlightNum + 1, cursor = Cursor + 1} + true -> + case next_pending_data(Tid, Cursor) of + none -> + Buffer; + #north_data{id = Id, tuple = Tuple} = NorthData -> + true = ets:insert(Tid, NorthData#north_data{inflight = true}), + ReceiverPid = self(), + ReceiverPid ! {next_data, Id, Tuple}, + endpoint_timer:task(TimerPid, Id, fun() -> ReceiverPid ! {next_data, Id, Tuple} end), + Buffer#buffer{flight_num = FlightNum + 1, cursor = Id} + end end. -spec ack(Id :: integer(), Buffer :: #buffer{}) -> NBuffer :: #buffer{}. -ack(Id, Buffer = #buffer{timer_pid = TimerPid, acc_num = AccNum, flight_num = FlightNum}) when is_integer(Id) -> +ack(Id, Buffer = #buffer{timer_pid = TimerPid, tid = Tid, acc_num = AccNum, flight_num = FlightNum}) when is_integer(Id) -> endpoint_timer:ack(TimerPid, Id), - trigger_next(Buffer#buffer{acc_num = AccNum + 1, flight_num = FlightNum - 1}). + case ets:take(Tid, Id) of + [#north_data{inflight = true}] -> + trigger_next(Buffer#buffer{acc_num = AccNum + 1, flight_num = max(FlightNum - 1, 0)}); + [#north_data{}] -> + trigger_next(Buffer#buffer{acc_num = AccNum + 1}); + [] -> + Buffer + end. %% 获取当前统计信息 -spec stat(Buffer :: #buffer{}) -> map(). -stat(#buffer{acc_num = AccNum, tid = Tid}) -> +stat(#buffer{acc_num = AccNum, tid = Tid, flight_num = FlightNum}) -> + QueueNum = max(ets:info(Tid, size) - FlightNum, 0), #{ <<"acc_num">> => AccNum, - <<"queue_num">> => ets:info(Tid, size) + <<"queue_num">> => QueueNum, + <<"inflight_num">> => FlightNum }. -spec cleanup(Buffer :: #buffer{}) -> ok. @@ -100,6 +115,33 @@ cleanup(#buffer{timer_pid = TimerPid}) -> endpoint_timer:cleanup(TimerPid), ok. +-spec resize(Buffer :: #buffer{}, WindowSize :: integer()) -> #buffer{}. +resize(Buffer = #buffer{}, WindowSize) when is_integer(WindowSize), WindowSize > 0 -> + trigger_n(Buffer#buffer{window_size = WindowSize}). + %%%=================================================================== %%% Internal functions -%%%=================================================================== \ No newline at end of file +%%%=================================================================== + +-spec next_pending_data(ets:tid(), integer()) -> none | #north_data{}. +next_pending_data(Tid, Cursor) -> + Key = case Cursor of + 0 -> + ets:first(Tid); + _ -> + ets:next(Tid, Cursor) + end, + next_pending_data_by_key(Tid, Key). + +-spec next_pending_data_by_key(ets:tid(), '$end_of_table' | integer()) -> none | #north_data{}. +next_pending_data_by_key(_Tid, '$end_of_table') -> + none; +next_pending_data_by_key(Tid, Key) -> + case ets:lookup(Tid, Key) of + [#north_data{inflight = false} = NorthData] -> + NorthData; + [_] -> + next_pending_data_by_key(Tid, ets:next(Tid, Key)); + [] -> + next_pending_data_by_key(Tid, ets:next(Tid, Key)) + end. diff --git a/src/endpoint/endpoint_http.erl b/src/endpoint/endpoint_http.erl index 7a01d8b..8b85746 100644 --- a/src/endpoint/endpoint_http.erl +++ b/src/endpoint/endpoint_http.erl @@ -44,10 +44,10 @@ start_link(LocalName, Endpoint = #endpoint{config = #http_endpoint{}}) when is_a -spec(init(Args :: term()) -> {ok, State :: #state{}} | {ok, State :: #state{}, timeout() | hibernate} | {stop, Reason :: term()} | ignore). -init([Endpoint = #endpoint{matcher = Matcher}]) -> +init([Endpoint = #endpoint{matcher = Matcher, config = #http_endpoint{pool_size = PoolSize}}]) -> ok = iot_log:set_metadata(), endpoint_subscription:subscribe(Matcher, self()), - Buffer = endpoint_buffer:new(Endpoint, 10), + Buffer = endpoint_buffer:new(Endpoint, PoolSize), {ok, #state{endpoint = Endpoint, buffer = Buffer}}. %% @private @@ -80,6 +80,10 @@ handle_cast({forward, Metric}, State = #state{buffer = Buffer, endpoint = #endpo end, NBuffer = endpoint_buffer:append(Tuple, Buffer), {noreply, State#state{buffer = NBuffer}}; +handle_cast({reload, NEndpoint = #endpoint{matcher = NMatcher, config = #http_endpoint{pool_size = PoolSize}}}, State = #state{endpoint = #endpoint{matcher = Matcher}, buffer = Buffer}) -> + ensure_subscription(Matcher, NMatcher), + NBuffer = endpoint_buffer:resize(Buffer, PoolSize), + {noreply, State#state{endpoint = NEndpoint, buffer = NBuffer}}; handle_cast(cleanup, State = #state{buffer = Buffer}) -> endpoint_buffer:cleanup(Buffer), @@ -100,19 +104,15 @@ handle_info({next_data, Id, {Metric, Sign}}, State = #state{buffer = Buffer, end Headers = BaseHeaders ++ ExtraHeaders, case hackney:request(post, Url, Headers, Metric) of - {ok, 200, _, ClientRef} -> - {ok, RespBody} = hackney:body(ClientRef), - hackney:close(ClientRef), + {ok, HttpCode, _, ClientRef} when HttpCode >= 200, HttpCode < 300 -> + RespBody = read_response_body(ClientRef), logger:debug("[endpoint_http] url: ~p, response is: ~p", [Url, RespBody]), NBuffer = endpoint_buffer:ack(Id, Buffer), - {noreply, State#state{buffer = NBuffer}}; {ok, HttpCode, _, ClientRef} -> - {ok, RespBody} = hackney:body(ClientRef), - hackney:close(ClientRef), - logger:debug("[endpoint_http] url: ~p, http_code: ~p, response is: ~p", [Url, HttpCode, RespBody]), - NBuffer = endpoint_buffer:ack(Id, Buffer), - {noreply, State#state{buffer = NBuffer}}; + RespBody = read_response_body(ClientRef), + logger:warning("[endpoint_http] url: ~p, http_code: ~p, response is: ~p", [Url, HttpCode, RespBody]), + {noreply, State}; {error, Reason} -> logger:warning("[endpoint_http] url: ~p, get error: ~p", [Url, Reason]), {noreply, State} @@ -125,7 +125,8 @@ handle_info({next_data, Id, {Metric, Sign}}, State = #state{buffer = Buffer, end %% with Reason. The return value is ignored. -spec(terminate(Reason :: (normal | shutdown | {shutdown, term()} | term()), State :: #state{}) -> term()). -terminate(_Reason, _State = #state{}) -> +terminate(_Reason, #state{buffer = Buffer}) -> + endpoint_buffer:cleanup(Buffer), ok. %% @private @@ -139,3 +140,21 @@ code_change(_OldVsn, State = #state{}, _Extra) -> %%%=================================================================== %%% Internal functions %%%=================================================================== + +-spec ensure_subscription(binary(), binary()) -> ok. +ensure_subscription(Matcher, Matcher) -> + ok; +ensure_subscription(Matcher, NMatcher) -> + ok = endpoint_subscription:unsubscribe(Matcher, self()), + endpoint_subscription:subscribe(NMatcher, self()). + +-spec read_response_body(reference()) -> binary() | term(). +read_response_body(ClientRef) -> + case hackney:body(ClientRef) of + {ok, RespBody} -> + hackney:close(ClientRef), + RespBody; + {error, Reason} -> + hackney:close(ClientRef), + Reason + end. diff --git a/src/endpoint/endpoint_kafka.erl b/src/endpoint/endpoint_kafka.erl index 6481119..ef7bae1 100644 --- a/src/endpoint/endpoint_kafka.erl +++ b/src/endpoint/endpoint_kafka.erl @@ -85,7 +85,15 @@ handle_call(get_stat, _From, State = #state{buffer = Buffer}) -> {stop, Reason :: term(), NewState :: #state{}}). handle_cast({forward, Metric}, State = #state{buffer = Buffer}) -> NBuffer = endpoint_buffer:append(Metric, Buffer), - {noreply, State#state{buffer = NBuffer}}. + {noreply, State#state{buffer = NBuffer}}; +handle_cast({reload, NEndpoint = #endpoint{matcher = NMatcher}}, State = #state{endpoint = #endpoint{matcher = Matcher}, client_id = ClientId}) -> + ensure_subscription(Matcher, NMatcher), + stop_kafka_client(ClientId), + retry_connect(), + {noreply, State#state{endpoint = NEndpoint, client_pid = undefined, status = ?DISCONNECTED}}; +handle_cast(cleanup, State = #state{buffer = Buffer}) -> + endpoint_buffer:cleanup(Buffer), + {noreply, State}. %% @private %% @doc Handling all non call/cast messages @@ -132,16 +140,29 @@ handle_info({next_data, _Id, _Tuple}, State = #state{status = ?DISCONNECTED}) -> {noreply, State}; %% 发送数据到mqtt服务器 handle_info({next_data, Id, Metric}, State = #state{status = ?CONNECTED, client_pid = ClientPid, - endpoint = #endpoint{config = #kafka_endpoint{topic = Topic}}}) -> + endpoint = #endpoint{config = #kafka_endpoint{topic = Topic}}, client_id = ClientId}) -> ReceiverPid = self(), AckCb = fun(Partition, BaseOffset) -> logger:debug("[endpoint_kafka] ack partion: ~p, offset: ~p", [Partition, BaseOffset]), ReceiverPid ! {ack, Id} end, - _ = brod:produce_cb(ClientPid, Topic, random, <<>>, Metric, AckCb), - - {noreply, State}; + case catch brod:produce_cb(ClientPid, Topic, random, <<>>, Metric, AckCb) of + {ok, _CallRef} -> + {noreply, State}; + {ok, _CallRef, _ProducerPid} -> + {noreply, State}; + {error, Reason} -> + logger:warning("[endpoint_kafka] produce topic: ~p, get error: ~p", [Topic, Reason]), + stop_kafka_client(ClientId), + retry_connect(), + {noreply, State#state{client_pid = undefined, status = ?DISCONNECTED}}; + {'EXIT', Reason} -> + logger:warning("[endpoint_kafka] produce topic: ~p, exit with reason: ~p", [Topic, Reason]), + stop_kafka_client(ClientId), + retry_connect(), + {noreply, State#state{client_pid = undefined, status = ?DISCONNECTED}} + end; handle_info({ack, Id}, State = #state{buffer = Buffer}) -> NBuffer = endpoint_buffer:ack(Id, Buffer), @@ -164,8 +185,9 @@ handle_info(Info, State = #state{status = Status}) -> %% with Reason. The return value is ignored. -spec(terminate(Reason :: (normal | shutdown | {shutdown, term()} | term()), State :: #state{}) -> term()). -terminate(Reason, #state{endpoint = #endpoint{title = Title}, buffer = Buffer}) -> +terminate(Reason, #state{endpoint = #endpoint{title = Title}, buffer = Buffer, client_id = ClientId}) -> logger:debug("[endpoint_kafka] endpoint: ~p, terminate with reason: ~p", [Title, Reason]), + stop_kafka_client(ClientId), endpoint_buffer:cleanup(Buffer), ok. @@ -183,3 +205,15 @@ code_change(_OldVsn, State = #state{}, _Extra) -> retry_connect() -> erlang:start_timer(?RETRY_INTERVAL, self(), connect). + +-spec ensure_subscription(binary(), binary()) -> ok. +ensure_subscription(Matcher, Matcher) -> + ok; +ensure_subscription(Matcher, NMatcher) -> + ok = endpoint_subscription:unsubscribe(Matcher, self()), + endpoint_subscription:subscribe(NMatcher, self()). + +-spec stop_kafka_client(atom()) -> ok. +stop_kafka_client(ClientId) when is_atom(ClientId) -> + _ = catch brod:stop_client(ClientId), + ok. diff --git a/src/endpoint/endpoint_mqtt.erl b/src/endpoint/endpoint_mqtt.erl index 272851b..8f5acf4 100644 --- a/src/endpoint/endpoint_mqtt.erl +++ b/src/endpoint/endpoint_mqtt.erl @@ -54,7 +54,7 @@ start_link(LocalName, Endpoint = #endpoint{}) when is_atom(LocalName) -> %% process to initialize. init([Endpoint = #endpoint{matcher = Matcher}]) -> ok = iot_log:set_metadata(), - % erlang:process_flag(trap_exit, true), + erlang:process_flag(trap_exit, true), endpoint_subscription:subscribe(Matcher, self()), %% 创建转发器, 避免阻塞当前进程的创建,因此采用了延时初始化的机制 @@ -86,7 +86,15 @@ handle_call(get_stat, _From, State = #state{buffer = Buffer}) -> {stop, Reason :: term(), NewState :: #state{}}). handle_cast({forward, Metric}, State = #state{buffer = Buffer}) -> NBuffer = endpoint_buffer:append(Metric, Buffer), - {noreply, State#state{buffer = NBuffer}}. + {noreply, State#state{buffer = NBuffer}}; +handle_cast({reload, NEndpoint = #endpoint{matcher = NMatcher}}, State = #state{endpoint = #endpoint{matcher = Matcher}, conn_pid = ConnPid}) -> + ensure_subscription(Matcher, NMatcher), + stop_mqtt_conn(ConnPid), + schedule_reconnect(), + {noreply, State#state{endpoint = NEndpoint, conn_pid = undefined, inflight = #{}, status = ?DISCONNECTED}}; +handle_cast(cleanup, State = #state{buffer = Buffer}) -> + endpoint_buffer:cleanup(Buffer), + {noreply, State}. %% @private %% @doc Handling all non call/cast messages @@ -133,7 +141,7 @@ handle_info({timeout, _, create_postman}, State = #state{buffer = Buffer, status %% 离线时,忽略数据发送逻辑 handle_info({next_data, _Id, _Tuple}, State = #state{status = ?DISCONNECTED}) -> - {keep_state, State}; + {noreply, State}; %% 发送数据到mqtt服务器 handle_info({next_data, Id, Metric}, State = #state{status = ?CONNECTED, conn_pid = ConnPid, buffer = Buffer, inflight = InFlight, endpoint = #endpoint{config = #mqtt_endpoint{topic = Topic, qos = Qos}}}) -> @@ -147,13 +155,15 @@ handle_info({next_data, Id, Metric}, State = #state{status = ?CONNECTED, conn_pi {noreply, State#state{inflight = maps:put(PacketId, Id, InFlight)}}; {error, Reason} -> logger:warning("[endpoint_mqtt] send message to topic: ~p, get error: ~p", [Topic, Reason]), - {stop, Reason, State} + stop_mqtt_conn(ConnPid), + schedule_reconnect(), + {noreply, State#state{conn_pid = undefined, inflight = #{}, status = ?DISCONNECTED}} end; -handle_info({disconnected, ReasonCode, Properties}, State = #state{status = ?CONNECTED}) -> + handle_info({disconnected, ReasonCode, Properties}, State = #state{status = ?CONNECTED}) -> logger:debug("[endpoint_mqtt] Recv a DISONNECT packet - ReasonCode: ~p, Properties: ~p", [ReasonCode, Properties]), - erlang:start_timer(?RETRY_INTERVAL, self(), create_postman), - {noreply, State#state{conn_pid = undefined, status = ?DISCONNECTED}}; + schedule_reconnect(), + {noreply, State#state{conn_pid = undefined, inflight = #{}, status = ?DISCONNECTED}}; handle_info({publish, Message = #{packet_id := _PacketId, payload := Payload}}, State = #state{status = ?CONNECTED}) -> logger:debug("[endpoint_mqtt] Recv a publish packet: ~p, payload: ~p", [Message, Payload]), @@ -172,8 +182,8 @@ handle_info({puback, #{packet_id := PacketId}}, State = #state{status = ?CONNECT %% postman进程挂掉时,重新建立新的 handle_info({'EXIT', ConnPid, Reason}, State = #state{endpoint = #endpoint{title = Title}, conn_pid = ConnPid}) -> logger:warning("[endpoint_mqtt] endpoint: ~p, conn pid exit with reason: ~p", [Title, Reason]), - erlang:start_timer(?RETRY_INTERVAL, self(), create_postman), - {noreply, State#state{conn_pid = undefined, status = ?DISCONNECTED}}; + schedule_reconnect(), + {noreply, State#state{conn_pid = undefined, inflight = #{}, status = ?DISCONNECTED}}; handle_info(Info, State = #state{status = Status}) -> logger:warning("[endpoint_mqtt] unknown message: ~p, status: ~p", [Info, Status]), @@ -186,8 +196,9 @@ handle_info(Info, State = #state{status = Status}) -> %% with Reason. The return value is ignored. -spec(terminate(Reason :: (normal | shutdown | {shutdown, term()} | term()), State :: #state{}) -> term()). -terminate(Reason, #state{endpoint = #endpoint{title = Title}, buffer = Buffer}) -> +terminate(Reason, #state{endpoint = #endpoint{title = Title}, buffer = Buffer, conn_pid = ConnPid}) -> logger:debug("[endpoint_mqtt] endpoint: ~p, terminate with reason: ~p", [Title, Reason]), + stop_mqtt_conn(ConnPid), endpoint_buffer:cleanup(Buffer), ok. @@ -202,3 +213,21 @@ code_change(_OldVsn, State = #state{}, _Extra) -> %%%=================================================================== %%% Internal functions %%%=================================================================== + +-spec ensure_subscription(binary(), binary()) -> ok. +ensure_subscription(Matcher, Matcher) -> + ok; +ensure_subscription(Matcher, NMatcher) -> + ok = endpoint_subscription:unsubscribe(Matcher, self()), + endpoint_subscription:subscribe(NMatcher, self()). + +-spec schedule_reconnect() -> reference(). +schedule_reconnect() -> + erlang:start_timer(?RETRY_INTERVAL, self(), create_postman). + +-spec stop_mqtt_conn(undefined | pid()) -> ok. +stop_mqtt_conn(undefined) -> + ok; +stop_mqtt_conn(ConnPid) when is_pid(ConnPid) -> + _ = catch emqtt:stop(ConnPid), + ok. diff --git a/src/endpoint/endpoint_subscription.erl b/src/endpoint/endpoint_subscription.erl index 95387d7..0369af0 100644 --- a/src/endpoint/endpoint_subscription.erl +++ b/src/endpoint/endpoint_subscription.erl @@ -13,7 +13,7 @@ %% API -export([start_link/0]). --export([subscribe/2, publish/2, get_subscribers/0]). +-export([subscribe/2, unsubscribe/2, publish/2, get_subscribers/0]). -export([match_components/2, is_valid_components/1, of_components/1]). %% gen_server callbacks @@ -26,6 +26,7 @@ topic :: binary(), subscriber_pid :: pid(), components = [], + monitor_ref :: undefined | reference(), %% 优先级 %% 1. 完全匹配的topic优先级别最高 %% 2. 带 * 的订阅 @@ -45,6 +46,10 @@ 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 :: map()}. get_subscribers() -> gen_server:call(?SERVER, get_subscribers). @@ -89,14 +94,34 @@ handle_call({subscribe, Topic, SubscriberPid}, _From, State = #state{subscribers Components = of_components(Topic), case is_valid_components(Components) of true -> - Sub = #subscriber{topic = Topic, subscriber_pid = SubscriberPid, components = Components, order = order_num(Components)}, - %% 建立到SubscriberPid的monitor,进程退出需要清理订阅 - erlang:monitor(process, SubscriberPid), - - {reply, ok, State#state{subscribers = Subscribers ++ [Sub]}}; + case has_subscription(Topic, SubscriberPid, Subscribers) of + true -> + {reply, ok, State}; + false -> + %% 建立到SubscriberPid的monitor,进程退出需要清理订阅 + MonitorRef = erlang:monitor(process, SubscriberPid), + Sub = #subscriber{ + topic = Topic, + subscriber_pid = SubscriberPid, + components = Components, + monitor_ref = MonitorRef, + order = order_num(Components) + }, + {reply, ok, State#state{subscribers = [Sub | Subscribers]}} + end; false -> {reply, {error, <<"invalid topic name">>}, State} - end. + end; +handle_call({unsubscribe, Topic, SubscriberPid}, _From, State = #state{subscribers = Subscribers}) -> + {Removed, Reserved} = lists:partition(fun(#subscriber{topic = Topic0, subscriber_pid = SubscriberPid0}) -> + Topic =:= Topic0 andalso SubscriberPid =:= SubscriberPid0 + end, Subscribers), + lists:foreach(fun(#subscriber{monitor_ref = MonitorRef}) when is_reference(MonitorRef) -> + erlang:demonitor(MonitorRef, [flush]); + (_) -> + ok + end, Removed), + {reply, ok, State#state{subscribers = Reserved}}. %% @private %% @doc Handling cast messages @@ -154,18 +179,11 @@ code_change(_OldVsn, State = #state{}, _Extra) -> -spec match_subscribers(Subscribers :: [#subscriber{}], Topic :: binary()) -> [#subscriber{}]. match_subscribers(Subscribers, Topic) when is_list(Subscribers), is_binary(Topic) -> Components = of_components(Topic), - lists:foldl(fun(S = #subscriber{components = Components0, subscriber_pid = Pid0}, Acc) -> - case match_components(Components0, Components) andalso not contain_channel(Pid0, Acc) of - true -> - [S|Acc]; - false -> - Acc - end - end, [], Subscribers). - --spec contain_channel(Pid :: pid(), Subscribers :: list()) -> boolean(). -contain_channel(Pid, Subscribers) when is_pid(Pid), is_list(Subscribers) -> - lists:search(fun(#subscriber{subscriber_pid = Pid0}) -> Pid == Pid0 end, Subscribers) /= false. + Matched = lists:filter(fun(#subscriber{components = Components0}) -> + match_components(Components0, Components) + end, Subscribers), + Sorted = lists:sort(fun compare_subscriber/2, Matched), + dedupe_subscribers(Sorted). %% 开始对比订阅的topic和发布的topic的Components信息 %% *表示单级匹配,+表示多级匹配;+只能出现一次,并且只能在末尾 @@ -205,3 +223,30 @@ order_num([<<$+>>|_]) -> 3; order_num([_|Tail]) -> order_num(Tail). + +-spec has_subscription(binary(), pid(), [#subscriber{}]) -> boolean(). +has_subscription(Topic, SubscriberPid, Subscribers) -> + lists:any(fun(#subscriber{topic = Topic0, subscriber_pid = SubscriberPid0}) -> + Topic =:= Topic0 andalso SubscriberPid =:= SubscriberPid0 + end, Subscribers). + +-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).