fix buffer

src/endpoint/endpoint_buffer.erl

@@ -15,15 +15,22 @@
 -define(RETRY_INTERVAL, 5000).
 %% 最大重试次数，不包含首次发送
 -define(MAX_RETRY_TIMES, 3).
+%% 与 endpoint_outbox 的单条记录限制保持一致；fast path 也不能绕过该限制。
+-define(MAX_PAYLOAD_BYTES, 32 * 1024).
 
--export([new/2, append/2, trigger_next/1, trigger_n/1, handle_timeout/4, cleanup/1, recover_inflight/1, ack/2, stat/1, resize/2]).
+-export([new/2, append/2, append_only/2, trigger_next/1, trigger_n/1, handle_timeout/4, cleanup/1, recover_inflight/1, ack/2, stat/1, resize/2]).
 -export_type([buffer/0]).
 
+-type flight_source() :: outbox | memory.
+-type timer_entry() :: {reference(), binary(), non_neg_integer(), flight_source()}.
+
 -record(buffer, {
     endpoint :: #endpoint{},
     outbox :: endpoint_outbox:outbox(),
-    %% 当前待 ack 的数据及其重试定时器 #{Id => {TimerRef, Payload, RetryTimes}}
+    %% 当前待 ack 的数据及其重试定时器 #{Id => {TimerRef, Payload, RetryTimes, Source}}
-    timer_map = #{} :: #{integer() => {reference(), binary(), non_neg_integer()}},
+    timer_map = #{} :: #{integer() => timer_entry()},
+    %% 内存 fast path 使用负数 id，避免与 outbox 的正整数 seq 冲突。
+    next_memory_id = -1 :: integer(),
     %% 窗口大小，允许最大的未确认消息数
     window_size = 10,
     %% 未确认的消息数
@@ -46,21 +53,38 @@ new(Endpoint = #endpoint{id = Id}, WindowSize) when is_integer(WindowSize), Wind
 
 -spec append(Payload :: binary(), Buffer :: #buffer{}) -> NBuffer :: #buffer{}.
 append(Payload, Buffer = #buffer{outbox = Outbox, window_size = WindowSize, flight_num = FlightNum}) when is_binary(Payload) ->
-    case endpoint_outbox:append(Payload, Outbox) of
+    case validate_payload_size(Payload, Buffer) of
-        {ok, _Seq, NOutbox} ->
+        ok ->
-            NBuffer = Buffer#buffer{outbox = NOutbox},
+            case FlightNum < WindowSize andalso outbox_empty(Outbox) of
-            case FlightNum < WindowSize of
                 true ->
-                    trigger_next(NBuffer);
+                    dispatch_memory(Payload, Buffer);
                 false ->
+                    case append_to_outbox(Payload, Buffer) of
+                        {ok, NBuffer} ->
+                            trigger_next(NBuffer);
+                        {dropped, NBuffer} ->
+                            trigger_next(NBuffer);
+                        {error, NBuffer} ->
+                            NBuffer
+                    end
+            end;
+        error ->
+            Buffer
+    end.
+
+-spec append_only(Payload :: binary(), Buffer :: #buffer{}) -> NBuffer :: #buffer{}.
+append_only(Payload, Buffer = #buffer{}) when is_binary(Payload) ->
+    case validate_payload_size(Payload, Buffer) of
+        ok ->
+            case append_to_outbox(Payload, Buffer) of
+                {ok, NBuffer} ->
+                    NBuffer;
+                {dropped, NBuffer} ->
+                    NBuffer;
+                {error, NBuffer} ->
                     NBuffer
             end;
-        {dropped, capacity_reached, NOutbox} ->
+        error ->
-            logger:warning("[endpoint_buffer] outbox capacity reached, endpoint_id: ~p", [buffer_endpoint_id(Buffer)]),
-            Buffer#buffer{outbox = NOutbox};
-        {error, Reason} ->
-            logger:warning("[endpoint_buffer] append outbox failed, endpoint_id: ~p, reason: ~p",
-                [buffer_endpoint_id(Buffer), Reason]),
             Buffer
     end.
 
@@ -82,7 +106,7 @@ trigger_next(Buffer = #buffer{outbox = Outbox, flight_num = FlightNum, window_si
                 {ok, Id, Payload, NOutbox} ->
                     ReceiverPid = self(),
                     ReceiverPid ! {next_data, Id, Payload},
-                    schedule_retry(Id, Payload, 0, Buffer#buffer{outbox = NOutbox, flight_num = FlightNum + 1});
+                    schedule_retry(Id, Payload, 0, outbox, Buffer#buffer{outbox = NOutbox, flight_num = FlightNum + 1});
                 {error, Reason} ->
                     logger:warning("[endpoint_buffer] read next outbox failed, endpoint_id: ~p, reason: ~p",
                         [buffer_endpoint_id(Buffer), Reason]),
@@ -94,17 +118,21 @@ trigger_next(Buffer = #buffer{outbox = Outbox, flight_num = FlightNum, window_si
 handle_timeout(TimerRef, Id, _Payload, Buffer = #buffer{timer_map = TimerMap})
         when is_reference(TimerRef), is_integer(Id) ->
     case maps:take(Id, TimerMap) of
-        {{TimerRef, Payload, RetryTimes}, NTimerMap} when RetryTimes < ?MAX_RETRY_TIMES ->
+        {{TimerRef, Payload, RetryTimes, Source}, NTimerMap} when RetryTimes < ?MAX_RETRY_TIMES ->
             logger:warning("[endpoint_buffer] retry message, endpoint_id: ~p, id: ~p, retry: ~p/~p",
                 [buffer_endpoint_id(Buffer), Id, RetryTimes + 1, ?MAX_RETRY_TIMES]),
             ReceiverPid = self(),
             ReceiverPid ! {next_data, Id, Payload},
-            schedule_retry(Id, Payload, RetryTimes + 1, Buffer#buffer{timer_map = NTimerMap});
+            schedule_retry(Id, Payload, RetryTimes + 1, Source, Buffer#buffer{timer_map = NTimerMap});
-        {{TimerRef, Payload, RetryTimes}, NTimerMap} ->
+        {{TimerRef, Payload, RetryTimes, outbox}, NTimerMap} ->
             logger:warning("[endpoint_buffer] drop message after retries exhausted, endpoint_id: ~p, id: ~p, retries: ~p",
                 [buffer_endpoint_id(Buffer), Id, RetryTimes]),
-            drop_message(Id, Payload, RetryTimes, Buffer#buffer{timer_map = NTimerMap});
+            drop_outbox_message(Id, Payload, RetryTimes, Buffer#buffer{timer_map = NTimerMap});
-        {{_OtherTimerRef, _OtherPayload, _RetryTimes}, _NTimerMap} ->
+        {{TimerRef, _, RetryTimes, memory}, NTimerMap} ->
+            logger:warning("[endpoint_buffer] drop memory message after retries exhausted, endpoint_id: ~p, id: ~p, retries: ~p",
+                [buffer_endpoint_id(Buffer), Id, RetryTimes]),
+            drop_memory_message(Buffer#buffer{timer_map = NTimerMap});
+        {{_OtherTimerRef, _OtherPayload, _RetryTimes, _Source}, _NTimerMap} ->
             Buffer;
         error ->
             Buffer
@@ -113,7 +141,15 @@ handle_timeout(TimerRef, Id, _Payload, Buffer = #buffer{timer_map = TimerMap})
 -spec ack(Id :: integer(), Buffer :: #buffer{}) -> NBuffer :: #buffer{}.
 ack(Id, Buffer = #buffer{timer_map = TimerMap, outbox = Outbox, acc_num = AccNum, flight_num = FlightNum}) when is_integer(Id) ->
     case maps:take(Id, TimerMap) of
-        {{TimerRef, Payload, RetryTimes}, NTimerMap} ->
+        {{TimerRef, _Payload, _RetryTimes, memory}, NTimerMap} ->
+            _ = erlang:cancel_timer(TimerRef),
+            NBuffer = Buffer#buffer{
+                timer_map = NTimerMap,
+                acc_num = AccNum + 1,
+                flight_num = max(FlightNum - 1, 0)
+            },
+            trigger_next(NBuffer);
+        {{TimerRef, Payload, RetryTimes, outbox}, NTimerMap} ->
             _ = erlang:cancel_timer(TimerRef),
             case endpoint_outbox:ack(Id, Outbox) of
                 {ok, NOutbox} ->
@@ -127,7 +163,7 @@ ack(Id, Buffer = #buffer{timer_map = TimerMap, outbox = Outbox, acc_num = AccNum
                 {error, Reason} ->
                     logger:warning("[endpoint_buffer] ack outbox failed, endpoint_id: ~p, id: ~p, reason: ~p",
                         [buffer_endpoint_id(Buffer), Id, Reason]),
-                    schedule_retry(Id, Payload, RetryTimes, Buffer#buffer{timer_map = NTimerMap})
+                    schedule_retry(Id, Payload, RetryTimes, outbox, Buffer#buffer{timer_map = NTimerMap})
             end;
         error ->
             Buffer
@@ -135,32 +171,38 @@ ack(Id, Buffer = #buffer{timer_map = TimerMap, outbox = Outbox, acc_num = AccNum
 
 %% 获取当前统计信息
 -spec stat(Buffer :: #buffer{}) -> map().
-stat(#buffer{acc_num = AccNum, outbox = Outbox, flight_num = FlightNum}) ->
+stat(#buffer{acc_num = AccNum, outbox = Outbox, flight_num = FlightNum, timer_map = TimerMap}) ->
     OutboxStat = endpoint_outbox:stat(Outbox),
     WriteSeq = maps:get(write_seq, OutboxStat, 0),
     AckedSeq = maps:get(acked_seq, OutboxStat, 0),
-    QueueNum = max(WriteSeq - AckedSeq - FlightNum, 0),
+    OutboxFlightNum = count_inflight(outbox, TimerMap),
+    MemoryFlightNum = count_inflight(memory, TimerMap),
+    QueueNum = max(WriteSeq - AckedSeq - OutboxFlightNum, 0),
     OutboxStat#{
         <<"acc_num">> => AccNum,
         <<"queue_num">> => QueueNum,
-        <<"inflight_num">> => FlightNum
+        <<"inflight_num">> => FlightNum,
+        <<"outbox_inflight_num">> => OutboxFlightNum,
+        <<"memory_inflight_num">> => MemoryFlightNum
     }.
 
 -spec cleanup(Buffer :: #buffer{}) -> #buffer{}.
 cleanup(Buffer = #buffer{timer_map = TimerMap}) ->
     cancel_timers(TimerMap),
-    Buffer#buffer{timer_map = #{}}.
+    NBuffer0 = persist_memory_inflight(Buffer),
+    reset_reader_after_recover(NBuffer0#buffer{timer_map = #{}, flight_num = 0}).
 
 -spec recover_inflight(Buffer :: #buffer{}) -> #buffer{}.
-recover_inflight(Buffer = #buffer{outbox = Outbox, timer_map = TimerMap}) ->
+recover_inflight(Buffer = #buffer{timer_map = TimerMap}) ->
     cancel_timers(TimerMap),
-    case endpoint_outbox:reset_reader(Outbox) of
+    NBuffer0 = persist_memory_inflight(Buffer),
+    case endpoint_outbox:reset_reader(NBuffer0#buffer.outbox) of
         {ok, NOutbox} ->
-            Buffer#buffer{outbox = NOutbox, timer_map = #{}, flight_num = 0};
+            NBuffer0#buffer{outbox = NOutbox, timer_map = #{}, flight_num = 0};
         {error, Reason} ->
             logger:warning("[endpoint_buffer] recover inflight failed, endpoint_id: ~p, reason: ~p",
                 [buffer_endpoint_id(Buffer), Reason]),
-            Buffer#buffer{timer_map = #{}, flight_num = 0}
+            NBuffer0#buffer{timer_map = #{}, flight_num = 0}
     end.
 
 -spec resize(Buffer :: #buffer{}, WindowSize :: integer()) -> #buffer{}.
@@ -175,22 +217,61 @@ resize(Buffer = #buffer{}, WindowSize) when is_integer(WindowSize), WindowSize >
 buffer_endpoint_id(#buffer{endpoint = #endpoint{id = Id}}) ->
     Id.
 
--spec schedule_retry(integer(), binary(), non_neg_integer(), buffer()) -> buffer().
+-spec validate_payload_size(binary(), buffer()) -> ok | error.
-schedule_retry(Id, Payload, RetryTimes, Buffer = #buffer{timer_map = TimerMap})
+validate_payload_size(Payload, Buffer) ->
+    case byte_size(Payload) =< ?MAX_PAYLOAD_BYTES of
+        true ->
+            ok;
+        false ->
+            logger:warning("[endpoint_buffer] payload too large, endpoint_id: ~p, size: ~p, max_size: ~p",
+                [buffer_endpoint_id(Buffer), byte_size(Payload), ?MAX_PAYLOAD_BYTES]),
+            error
+    end.
+
+-spec outbox_empty(endpoint_outbox:outbox()) -> boolean().
+outbox_empty(Outbox) ->
+    Stat = endpoint_outbox:stat(Outbox),
+    maps:get(write_seq, Stat, 0) =:= maps:get(acked_seq, Stat, 0).
+
+-spec append_to_outbox(binary(), buffer()) -> {ok | dropped | error, buffer()}.
+append_to_outbox(Payload, Buffer = #buffer{outbox = Outbox}) ->
+    case endpoint_outbox:append(Payload, Outbox) of
+        {ok, _Seq, NOutbox} ->
+            {ok, Buffer#buffer{outbox = NOutbox}};
+        {dropped, capacity_reached, NOutbox} ->
+            logger:warning("[endpoint_buffer] outbox capacity reached, endpoint_id: ~p", [buffer_endpoint_id(Buffer)]),
+            {dropped, Buffer#buffer{outbox = NOutbox}};
+        {error, Reason} ->
+            logger:warning("[endpoint_buffer] append outbox failed, endpoint_id: ~p, reason: ~p",
+                [buffer_endpoint_id(Buffer), Reason]),
+            {error, Buffer}
+    end.
+
+-spec dispatch_memory(binary(), buffer()) -> buffer().
+dispatch_memory(Payload, Buffer = #buffer{next_memory_id = Id, flight_num = FlightNum}) ->
+    ReceiverPid = self(),
+    ReceiverPid ! {next_data, Id, Payload},
+    schedule_retry(Id, Payload, 0, memory, Buffer#buffer{
+        next_memory_id = Id - 1,
+        flight_num = FlightNum + 1
+    }).
+
+-spec schedule_retry(integer(), binary(), non_neg_integer(), flight_source(), buffer()) -> buffer().
+schedule_retry(Id, Payload, RetryTimes, Source, Buffer = #buffer{timer_map = TimerMap})
         when is_integer(Id), is_binary(Payload), is_integer(RetryTimes), RetryTimes >= 0 ->
     TimerRef = erlang:start_timer(?RETRY_INTERVAL, self(), {endpoint_buffer_retry, Id, Payload}),
-    Buffer#buffer{timer_map = maps:put(Id, {TimerRef, Payload, RetryTimes}, TimerMap)}.
+    Buffer#buffer{timer_map = maps:put(Id, {TimerRef, Payload, RetryTimes, Source}, TimerMap)}.
 
--spec cancel_timers(#{integer() => {reference(), binary(), non_neg_integer()}}) -> ok.
+-spec cancel_timers(#{integer() => timer_entry()}) -> ok.
 cancel_timers(TimerMap) ->
-    lists:foreach(fun({_Id, {TimerRef, _Payload, _RetryTimes}}) ->
+    lists:foreach(fun({_Id, {TimerRef, _Payload, _RetryTimes, _Source}}) ->
         _ = erlang:cancel_timer(TimerRef),
         ok
     end, maps:to_list(TimerMap)),
     ok.
 
--spec drop_message(integer(), binary(), non_neg_integer(), buffer()) -> buffer().
+-spec drop_outbox_message(integer(), binary(), non_neg_integer(), buffer()) -> buffer().
-drop_message(Id, Payload, RetryTimes, Buffer = #buffer{outbox = Outbox, acc_num = AccNum, flight_num = FlightNum}) ->
+drop_outbox_message(Id, Payload, RetryTimes, Buffer = #buffer{outbox = Outbox, acc_num = AccNum, flight_num = FlightNum}) ->
     case endpoint_outbox:ack(Id, Outbox) of
         {ok, NOutbox} ->
             trigger_next(Buffer#buffer{
@@ -201,5 +282,51 @@ drop_message(Id, Payload, RetryTimes, Buffer = #buffer{outbox = Outbox, acc_num
         {error, Reason} ->
             logger:warning("[endpoint_buffer] ack dropped message failed, endpoint_id: ~p, id: ~p, reason: ~p",
                 [buffer_endpoint_id(Buffer), Id, Reason]),
-            schedule_retry(Id, Payload, RetryTimes, Buffer)
+            schedule_retry(Id, Payload, RetryTimes, outbox, Buffer)
     end.
+
+-spec drop_memory_message(buffer()) -> buffer().
+drop_memory_message(Buffer = #buffer{acc_num = AccNum, flight_num = FlightNum}) ->
+    trigger_next(Buffer#buffer{
+        acc_num = AccNum + 1,
+        flight_num = max(FlightNum - 1, 0)
+    }).
+
+-spec persist_memory_inflight(buffer()) -> buffer().
+persist_memory_inflight(Buffer = #buffer{timer_map = TimerMap}) ->
+    MemoryInflight = lists:sort(
+        fun({IdA, _PayloadA}, {IdB, _PayloadB}) -> IdA > IdB end,
+        [{Id, Payload} || {Id, {_TimerRef, Payload, _RetryTimes, memory}} <- maps:to_list(TimerMap)]
+    ),
+    lists:foldl(fun({_Id, Payload}, AccBuffer) ->
+        case append_to_outbox(Payload, AccBuffer) of
+            {ok, NBuffer} ->
+                NBuffer;
+            {dropped, NBuffer} ->
+                NBuffer;
+            {error, NBuffer} ->
+                NBuffer
+        end
+    end, Buffer, MemoryInflight).
+
+-spec reset_reader_after_recover(buffer()) -> buffer().
+reset_reader_after_recover(Buffer = #buffer{outbox = Outbox}) ->
+    case endpoint_outbox:reset_reader(Outbox) of
+        {ok, NOutbox} ->
+            Buffer#buffer{outbox = NOutbox};
+        {error, Reason} ->
+            logger:warning("[endpoint_buffer] reset reader failed, endpoint_id: ~p, reason: ~p",
+                [buffer_endpoint_id(Buffer), Reason]),
+            Buffer
+    end.
+
+-spec count_inflight(flight_source(), #{integer() => timer_entry()}) -> non_neg_integer().
+count_inflight(Source, TimerMap) ->
+    maps:fold(fun(_Id, {_TimerRef, _Payload, _RetryTimes, EntrySource}, Acc) ->
+        case EntrySource =:= Source of
+            true ->
+                Acc + 1;
+            false ->
+                Acc
+        end
+    end, 0, TimerMap).

src/endpoint/endpoint_kafka.erl

@@ -62,7 +62,7 @@ init([Endpoint = #endpoint{id = Id, matcher = Matcher}]) ->
 disconnected({call, From}, get_stat, State = #state{buffer = Buffer}) ->
     reply_stat(From, Buffer, State);
 disconnected(cast, {forward, Metric}, State = #state{buffer = Buffer}) ->
-    forward_metric(Metric, Buffer, State);
+    store_metric(Metric, Buffer, State);
 disconnected(cast, cleanup, State = #state{buffer = Buffer}) ->
     cleanup_buffer(Buffer, State);
 disconnected(cast, {reload, NEndpoint = #endpoint{matcher = NMatcher}},
@@ -167,6 +167,12 @@ forward_metric(Metric, Buffer, State) ->
     NBuffer = endpoint_buffer:append(Metric, Buffer),
     {keep_state, State#state{buffer = NBuffer}}.
 
+-spec store_metric(binary(), endpoint_buffer:buffer(), #state{}) ->
+    gen_statem:event_handler_result(kafka_state(), #state{}).
+store_metric(Metric, Buffer, State) ->
+    NBuffer = endpoint_buffer:append_only(Metric, Buffer),
+    {keep_state, State#state{buffer = NBuffer}}.
+
 -spec cleanup_buffer(endpoint_buffer:buffer(), #state{}) ->
     gen_statem:event_handler_result(kafka_state(), #state{}).
 cleanup_buffer(Buffer, _State) ->

src/endpoint/endpoint_mqtt.erl

@@ -61,7 +61,7 @@ init([Endpoint = #endpoint{matcher = Matcher}]) ->
 disconnected({call, From}, get_stat, State = #state{buffer = Buffer}) ->
     reply_stat(From, Buffer, State);
 disconnected(cast, {forward, Metric}, State = #state{buffer = Buffer}) ->
-    forward_metric(Metric, Buffer, State);
+    store_metric(Metric, Buffer, State);
 disconnected(cast, cleanup, State = #state{buffer = Buffer}) ->
     cleanup_buffer(Buffer, State);
 disconnected(cast, {reload, NEndpoint = #endpoint{matcher = NMatcher}},
@@ -220,6 +220,12 @@ forward_metric(Metric, Buffer, State) ->
     NBuffer = endpoint_buffer:append(Metric, Buffer),
     {keep_state, State#state{buffer = NBuffer}}.
 
+-spec store_metric(binary(), endpoint_buffer:buffer(), #state{}) ->
+    gen_statem:event_handler_result(mqtt_state(), #state{}).
+store_metric(Metric, Buffer, State) ->
+    NBuffer = endpoint_buffer:append_only(Metric, Buffer),
+    {keep_state, State#state{buffer = NBuffer}}.
+
 -spec cleanup_buffer(endpoint_buffer:buffer(), #state{}) ->
     gen_statem:event_handler_result(mqtt_state(), #state{}).
 cleanup_buffer(Buffer, State) ->