fix endpoint

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{
+    Errors = lists:filtermap(fun(Term) ->
-        url = Url,
+        case check_http_argument(Term) of
-        token = Token,
+            ok ->
-        pool_size = PoolSize
+                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

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}) ->
+trigger_next(Buffer = #buffer{tid = Tid, cursor = Cursor, timer_pid = TimerPid, flight_num = FlightNum, window_size = WindowSize}) ->
-    case ets:first(Tid) of
+    case FlightNum < WindowSize of
-        '$end_of_table' ->
+        false ->
             Buffer;
-        Key ->
+        true ->
-            [#north_data{id = Id, tuple = Tuple}|_] = ets:take(Tid, Key),
+            case next_pending_data(Tid, Cursor) of
-            ReceiverPid = self(),
+                none ->
-            ReceiverPid ! {next_data, Id, Tuple},
+                    Buffer;
-            endpoint_timer:task(TimerPid, Id, fun() -> ReceiverPid ! {next_data, Id, Tuple} end),
+                #north_data{id = Id, tuple = Tuple} = NorthData ->
-            Buffer#buffer{flight_num = FlightNum + 1, cursor = Cursor + 1}
+                    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
 %%%===================================================================
+
+-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.

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, HttpCode, _, ClientRef} when HttpCode >= 200, HttpCode < 300 ->
-            {ok, RespBody} = hackney:body(ClientRef),
+            RespBody = read_response_body(ClientRef),
-            hackney:close(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),
+            RespBody = read_response_body(ClientRef),
-            hackney:close(ClientRef),
+            logger:warning("[endpoint_http] url: ~p, http_code: ~p, response is: ~p", [Url, HttpCode, RespBody]),
-            logger:debug("[endpoint_http] url: ~p, http_code: ~p, response is: ~p", [Url, HttpCode, RespBody]),
+            {noreply, State};
-            NBuffer = endpoint_buffer:ack(Id, Buffer),
-            {noreply, State#state{buffer = NBuffer}};
         {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.

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),
+    case catch brod:produce_cb(ClientPid, Topic, random, <<>>, Metric, AckCb) of
-
+        {ok, _CallRef} ->
-    {noreply, State};
+            {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.

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),
+    schedule_reconnect(),
-    {noreply, State#state{conn_pid = undefined, status = ?DISCONNECTED}};
+    {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),
+    schedule_reconnect(),
-    {noreply, State#state{conn_pid = undefined, status = ?DISCONNECTED}};
+    {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.

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)},
+            case has_subscription(Topic, SubscriberPid, Subscribers) of
-            %% 建立到SubscriberPid的monitor，进程退出需要清理订阅
+                true ->
-            erlang:monitor(process, SubscriberPid),
+                    {reply, ok, State};
-
+                false ->
-            {reply, ok, State#state{subscribers = Subscribers ++ [Sub]}};
+                    %% 建立到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) ->
+    Matched = lists:filter(fun(#subscriber{components = Components0}) ->
-        case match_components(Components0, Components) andalso not contain_channel(Pid0, Acc) of
+        match_components(Components0, Components)
-            true ->
+    end, Subscribers),
-                [S|Acc];
+    Sorted = lists:sort(fun compare_subscriber/2, Matched),
-            false ->
+    dedupe_subscribers(Sorted).
-                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.
 
 %% 开始对比订阅的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).