diff --git a/src/endpoint/endpoint_buffer.erl b/src/endpoint/endpoint_buffer.erl index ac5cc02..d69ce51 100644 --- a/src/endpoint/endpoint_buffer.erl +++ b/src/endpoint/endpoint_buffer.erl @@ -14,7 +14,7 @@ %% 消息重发间隔 -define(RETRY_INTERVAL, 5000). --export([new/2, append/2, trigger_next/1, trigger_n/1, cleanup/1, ack/2, stat/1, resize/2]). +-export([new/2, append/2, trigger_next/1, trigger_n/1, cleanup/1, recover_inflight/1, ack/2, stat/1, resize/2]). -export_type([buffer/0]). -record(buffer, { @@ -110,10 +110,16 @@ stat(#buffer{acc_num = AccNum, tid = Tid, flight_num = FlightNum}) -> <<"inflight_num">> => FlightNum }. --spec cleanup(Buffer :: #buffer{}) -> ok. -cleanup(#buffer{timer_pid = TimerPid}) -> +-spec cleanup(Buffer :: #buffer{}) -> #buffer{}. +cleanup(Buffer = #buffer{timer_pid = TimerPid}) -> endpoint_timer:cleanup(TimerPid), - ok. + Buffer. + +-spec recover_inflight(Buffer :: #buffer{}) -> #buffer{}. +recover_inflight(Buffer = #buffer{tid = Tid, timer_pid = TimerPid}) -> + endpoint_timer:cleanup(TimerPid), + reset_inflight(Tid, ets:first(Tid)), + Buffer#buffer{cursor = 0, flight_num = 0}. -spec resize(Buffer :: #buffer{}, WindowSize :: integer()) -> #buffer{}. resize(Buffer = #buffer{}, WindowSize) when is_integer(WindowSize), WindowSize > 0 -> @@ -145,3 +151,18 @@ next_pending_data_by_key(Tid, Key) -> [] -> next_pending_data_by_key(Tid, ets:next(Tid, Key)) end. + +-spec reset_inflight(ets:tid(), '$end_of_table' | integer()) -> ok. +reset_inflight(_Tid, '$end_of_table') -> + ok; +reset_inflight(Tid, Key) -> + NextKey = ets:next(Tid, Key), + case ets:lookup(Tid, Key) of + [#north_data{inflight = true} = NorthData] -> + true = ets:insert(Tid, NorthData#north_data{inflight = false}); + [_] -> + ok; + [] -> + ok + end, + reset_inflight(Tid, NextKey). diff --git a/src/endpoint/endpoint_http.erl b/src/endpoint/endpoint_http.erl index 8b85746..5bc991a 100644 --- a/src/endpoint/endpoint_http.erl +++ b/src/endpoint/endpoint_http.erl @@ -86,8 +86,8 @@ handle_cast({reload, NEndpoint = #endpoint{matcher = NMatcher, config = #http_en {noreply, State#state{endpoint = NEndpoint, buffer = NBuffer}}; handle_cast(cleanup, State = #state{buffer = Buffer}) -> - endpoint_buffer:cleanup(Buffer), - {noreply, State}. + NBuffer = endpoint_buffer:cleanup(Buffer), + {noreply, State#state{buffer = NBuffer}}. %% @private %% @doc Handling all non call/cast messages @@ -126,7 +126,7 @@ handle_info({next_data, Id, {Metric, Sign}}, State = #state{buffer = Buffer, end -spec(terminate(Reason :: (normal | shutdown | {shutdown, term()} | term()), State :: #state{}) -> term()). terminate(_Reason, #state{buffer = Buffer}) -> - endpoint_buffer:cleanup(Buffer), + _ = endpoint_buffer:cleanup(Buffer), ok. %% @private diff --git a/src/endpoint/endpoint_kafka.erl b/src/endpoint/endpoint_kafka.erl index 35264c5..08c5c9d 100644 --- a/src/endpoint/endpoint_kafka.erl +++ b/src/endpoint/endpoint_kafka.erl @@ -92,6 +92,7 @@ connected(cast, {reload, NEndpoint = #endpoint{matcher = NMatcher}}, reload_endpoint(Matcher, NMatcher, ClientId, NEndpoint, State); connected(info, {next_data, Id, Metric}, State = #state{ + buffer = Buffer, client_pid = ClientPid, client_id = ClientId, endpoint = #endpoint{config = #kafka_endpoint{topic = Topic}} @@ -109,20 +110,23 @@ connected(info, {next_data, Id, Metric}, {error, Reason} -> logger:warning("[endpoint_kafka] produce topic: ~p, get error: ~p", [Topic, Reason]), stop_kafka_client(ClientId), - {next_state, disconnected, State#state{client_pid = undefined}, + NBuffer = endpoint_buffer:recover_inflight(Buffer), + {next_state, disconnected, State#state{client_pid = undefined, buffer = NBuffer}, [{state_timeout, ?RETRY_INTERVAL, connect}]}; {'EXIT', Reason} -> logger:warning("[endpoint_kafka] produce topic: ~p, exit with reason: ~p", [Topic, Reason]), stop_kafka_client(ClientId), - {next_state, disconnected, State#state{client_pid = undefined}, + NBuffer = endpoint_buffer:recover_inflight(Buffer), + {next_state, disconnected, State#state{client_pid = undefined, buffer = NBuffer}, [{state_timeout, ?RETRY_INTERVAL, connect}]} end; connected(info, {ack, Id}, State = #state{buffer = Buffer}) -> ack_buffer(Id, Buffer, State); connected(info, {'EXIT', ClientPid, Reason}, - State = #state{client_pid = ClientPid, endpoint = #endpoint{title = Title}}) -> + State = #state{client_pid = ClientPid, endpoint = #endpoint{title = Title}, buffer = Buffer}) -> logger:warning("[endpoint_kafka] endpoint: ~p, conn pid exit with reason: ~p", [Title, Reason]), - {next_state, disconnected, State#state{client_pid = undefined}, + NBuffer = endpoint_buffer:recover_inflight(Buffer), + {next_state, disconnected, State#state{client_pid = undefined, buffer = NBuffer}, [{state_timeout, ?RETRY_INTERVAL, connect}]}; connected(info, Info, State) -> unknown_info(Info, connected, State); @@ -160,8 +164,8 @@ forward_metric(Metric, Buffer, State) -> -spec cleanup_buffer(endpoint_buffer:buffer(), #state{}) -> gen_statem:event_handler_result(kafka_state(), #state{}). cleanup_buffer(Buffer, _State) -> - endpoint_buffer:cleanup(Buffer), - keep_state_and_data. + NBuffer = endpoint_buffer:cleanup(Buffer), + {keep_state, _State#state{buffer = NBuffer}}. -spec ack_buffer(integer(), endpoint_buffer:buffer(), #state{}) -> gen_statem:event_handler_result(kafka_state(), #state{}). @@ -174,7 +178,8 @@ ack_buffer(Id, Buffer, State) -> reload_endpoint(Matcher, NMatcher, ClientId, NEndpoint, State = #state{}) -> ensure_subscription(Matcher, NMatcher), stop_kafka_client(ClientId), - {next_state, disconnected, State#state{endpoint = NEndpoint, client_pid = undefined}, + NBuffer = endpoint_buffer:recover_inflight(State#state.buffer), + {next_state, disconnected, State#state{endpoint = NEndpoint, client_pid = undefined, buffer = NBuffer}, [{state_timeout, 0, connect}]}. -spec try_connect(#state{}) -> gen_statem:event_handler_result(kafka_state(), #state{}). diff --git a/src/endpoint/endpoint_mqtt.erl b/src/endpoint/endpoint_mqtt.erl index de6b4aa..d5ee42b 100644 --- a/src/endpoint/endpoint_mqtt.erl +++ b/src/endpoint/endpoint_mqtt.erl @@ -152,13 +152,16 @@ connected(info, {next_data, Id, Metric}, {error, Reason} -> logger:warning("[endpoint_mqtt] send message to topic: ~p, get error: ~p", [Topic, Reason]), stop_mqtt_conn(ConnPid), - {next_state, disconnected, State#state{conn_pid = undefined, inflight = #{}}, + NBuffer = endpoint_buffer:recover_inflight(Buffer), + {next_state, disconnected, State#state{conn_pid = undefined, inflight = #{}, buffer = NBuffer}, [{state_timeout, ?RETRY_INTERVAL, connect}]} end; -connected(info, {disconnected, ReasonCode, Properties}, State = #state{conn_pid = ConnPid}) -> +connected(info, {disconnected, ReasonCode, Properties}, + State = #state{conn_pid = ConnPid, buffer = Buffer}) -> logger:debug("[endpoint_mqtt] Recv a DISONNECT packet - ReasonCode: ~p, Properties: ~p", [ReasonCode, Properties]), stop_mqtt_conn(ConnPid), - {next_state, disconnected, State#state{conn_pid = undefined, inflight = #{}}, + NBuffer = endpoint_buffer:recover_inflight(Buffer), + {next_state, disconnected, State#state{conn_pid = undefined, inflight = #{}, buffer = NBuffer}, [{state_timeout, ?RETRY_INTERVAL, connect}]}; connected(info, {publish, Message = #{packet_id := _PacketId, payload := Payload}}, State) -> logger:debug("[endpoint_mqtt] Recv a publish packet: ~p, payload: ~p", [Message, Payload]), @@ -173,9 +176,10 @@ connected(info, {puback, #{packet_id := PacketId}}, {keep_state, State} end; connected(info, {'EXIT', ConnPid, Reason}, - State = #state{endpoint = #endpoint{title = Title}, conn_pid = ConnPid}) -> + State = #state{endpoint = #endpoint{title = Title}, conn_pid = ConnPid, buffer = Buffer}) -> logger:warning("[endpoint_mqtt] endpoint: ~p, conn pid exit with reason: ~p", [Title, Reason]), - {next_state, disconnected, State#state{conn_pid = undefined, inflight = #{}}, + NBuffer = endpoint_buffer:recover_inflight(Buffer), + {next_state, disconnected, State#state{conn_pid = undefined, inflight = #{}, buffer = NBuffer}, [{state_timeout, ?RETRY_INTERVAL, connect}]}; connected(info, Info, State) -> unknown_info(Info, connected, State); @@ -212,16 +216,17 @@ forward_metric(Metric, Buffer, State) -> -spec cleanup_buffer(endpoint_buffer:buffer(), #state{}) -> gen_statem:event_handler_result(mqtt_state(), #state{}). -cleanup_buffer(Buffer, _State) -> - endpoint_buffer:cleanup(Buffer), - keep_state_and_data. +cleanup_buffer(Buffer, State) -> + NBuffer = endpoint_buffer:cleanup(Buffer), + {keep_state, State#state{buffer = NBuffer}}. -spec reload_endpoint(binary(), binary(), undefined | pid(), #endpoint{}, #state{}) -> gen_statem:event_handler_result(mqtt_state(), #state{}). reload_endpoint(Matcher, NMatcher, ConnPid, NEndpoint, State = #state{}) -> ensure_subscription(Matcher, NMatcher), stop_mqtt_conn(ConnPid), - {next_state, disconnected, State#state{endpoint = NEndpoint, conn_pid = undefined, inflight = #{}}, + NBuffer = endpoint_buffer:recover_inflight(State#state.buffer), + {next_state, disconnected, State#state{endpoint = NEndpoint, conn_pid = undefined, inflight = #{}, buffer = NBuffer}, [{next_event, internal, do_connect}]}. -spec unknown_info(term(), mqtt_state(), #state{}) ->