Add per-node ingress byte tracking to rabbit_fifo v9

kjnilsson · kjnilsson · commit 07e347884e54 · 2026-06-02T14:00:00.000+01:00
Implement the OSS infrastructure for ingress-dependent quorum queue leader
rebalancing. This adds:

- Per-node enqueue byte counters in the replicated state machine, updated
  deterministically on every enqueue via apply_enqueue/6
- Exposure of cumulative totals via overview/1, propagated to all replicas
  on every tick
- Aux state leaky integrators (via ra_li) to compute smoothed per-node
  ingress rates, updated on handle_tick
- A get_ingress_rates aux query that returns bytes/second per node,
  callable locally without cross-node coordination

The new field ingress_bytes_by_node defaults to #{} and is never read by
OSS, so pre-v9 members in mixed clusters are unaffected. The aux state
version is bumped to aux_v5 with an upgrade clause for rolling restarts.

Decay time is configurable via persistent_term.

Tests added to verify accumulation, snapshot persistence, v8→v9 conversion,
and rate decay behavior under controlled tick timing.
diff --git a/deps/rabbit/src/rabbit_fifo.erl b/deps/rabbit/src/rabbit_fifo.erl
@@ -203,10 +203,15 @@
               delayed_op/0]).
 
 -spec init(config()) -> state().
-init(#{name := Name,
-       queue_resource := Resource} = Conf) ->
-    update_config(Conf, #?STATE{cfg = #cfg{name = Name,
-                                           resource = Resource}}).
+init(Conf) ->
+    case Conf of
+        #{name := Name, queue_resource := Resource} ->
+            update_config(Conf, #?STATE{cfg = #cfg{name = Name,
+                                                   resource = Resource}});
+        _ ->
+            erlang:display({bad_init_arg, Conf}),
+            erlang:error({bad_init_arg, Conf})
+    end.
 
 update_config(Conf, State) ->
     DLH = maps:get(dead_letter_handler, Conf, undefined),
@@ -696,9 +701,12 @@ apply_(Meta, {nodeup, Node}, #?STATE{consumers = Cons0,
 apply_(_, {nodedown, _Node}, State) ->
     {State, ok};
 apply_(Meta, #purge_nodes{nodes = Nodes}, State0) ->
-    {State, Effects} = lists:foldl(fun(Node, {S, E}) ->
+    {State1, Effects} = lists:foldl(fun(Node, {S, E}) ->
                                            purge_node(Meta, Node, S, E)
                                    end, {State0, []}, Nodes),
+    State = State1#?STATE{
+        ingress_bytes_by_node =
+            maps:without(Nodes, State1#?STATE.ingress_bytes_by_node)},
     {State, ok, Effects};
 apply_(Meta,
        #update_config{config = #{} = Conf},
@@ -967,8 +975,7 @@ credit_reply_resend_effect(#?STATE{waiting_consumers = Waiting,
       end, [], maps:merge(Consumers, maps:from_list(Waiting))).
 
 convert_v8_to_v9(#{} = _Meta, StateV8) ->
-    State = StateV8,
-    State.
+    erlang:append_element(StateV8, #{}).
 
 purge_node(Meta, Node, State, Effects) ->
     lists:foldl(fun(Pid, {S0, E0}) ->
@@ -1171,7 +1178,8 @@ overview(#?STATE{consumers = Cons,
                  reclaimable_bytes_count => ReclaimableBytes,
                  smallest_raft_index => smallest_raft_index(State),
                  num_active_priorities => NumActivePriorities,
-                 messages_by_priority => Detail
+                 messages_by_priority => Detail,
+                 ingress_bytes_by_node => State#?STATE.ingress_bytes_by_node
                  },
     DlxOverview = dlx_overview(DlxState),
     maps:merge(maps:merge(Overview, DlxOverview), SacOverview).
@@ -1205,7 +1213,13 @@ which_module(7) -> rabbit_fifo_v7;
 which_module(8) -> rabbit_fifo_v8;
 which_module(9) -> ?MODULE.
 
--define(AUX, aux_v4).
+-define(AUX, aux_v5).
+-define(DEFAULT_INGRESS_DECAY_MS, 60_000).
+
+-record(ingress_aux,
+        {last_totals = #{} :: #{node() | undefined => non_neg_integer()},
+         estimators = #{} :: #{node() | undefined => ra_li:state()},
+         decay_ms = ?DEFAULT_INGRESS_DECAY_MS :: pos_integer()}).
 
 -record(snapshot, {index :: ra:index(),
                    timestamp :: milliseconds(),
@@ -1220,7 +1234,8 @@ which_module(9) -> ?MODULE.
                gc = #aux_gc{} :: #aux_gc{},
                tick_pid :: undefined | pid(),
                cache = #{} :: map(),
-               last_checkpoint :: tuple() | #snapshot{}
+               last_checkpoint :: tuple() | #snapshot{},
+               ingress = #ingress_aux{} :: #ingress_aux{}
               }).
 
 init_aux(Name) when is_atom(Name) ->
@@ -1234,11 +1249,38 @@ init_aux(Name) when is_atom(Name) ->
                              ?SNAP_MIN_RECLAIMABLE_B}),
     Range = max(1, SnapMinReclaimable - ?SNAP_MIN_RECLAIMABLE_LOW_B),
     MinReclaimable = ?SNAP_MIN_RECLAIMABLE_LOW_B + rand:uniform(Range),
+    DecayMs = persistent_term:get(rabbit_fifo_ingress_decay_ms,
+                                  ?DEFAULT_INGRESS_DECAY_MS),
     #?AUX{name = Name,
           last_checkpoint = #snapshot{index = 0,
                                       timestamp = erlang:system_time(millisecond),
                                       messages_total = 0,
-                                      min_reclaimable = MinReclaimable}}.
+                                      min_reclaimable = MinReclaimable},
+          ingress = #ingress_aux{decay_ms = DecayMs}}.
+
+update_ingress(Overview, Nodes, #ingress_aux{last_totals = LastTotals,
+                                              estimators = Estimators0,
+                                              decay_ms = DecayMs} = Ingress) ->
+    NewTotals = maps:get(ingress_bytes_by_node, Overview, #{}),
+    Ts = erlang:monotonic_time(millisecond),
+    Estimators1 =
+        maps:fold(fun(Node, NewTotal, Est) ->
+                      Delta = NewTotal - maps:get(Node, LastTotals, 0),
+                      Li0 = maps:get(Node, Est, ra_li:new(DecayMs)),
+                      Li1 = ra_li:update(Delta, Ts, Li0),
+                      Est#{Node => Li1}
+                  end, Estimators0, NewTotals),
+    ActiveNodes = sets:from_list(Nodes, [{version, 2}]),
+    Estimators = maps:filter(fun(Node, _) ->
+                                  Node =:= undefined orelse
+                                  sets:is_element(Node, ActiveNodes)
+                             end, Estimators1),
+    Ingress#ingress_aux{last_totals = NewTotals,
+                        estimators = Estimators}.
+
+compute_ingress_rates(#ingress_aux{estimators = Estimators}) ->
+    Ts = erlang:monotonic_time(millisecond),
+    maps:map(fun(_Node, Li) -> ra_li:rate(Ts, Li) end, Estimators).
 
 handle_aux(RaftState, Tag, Cmd, AuxV2, RaAux)
   when element(1, AuxV2) == aux_v2 ->
@@ -1256,6 +1298,19 @@ handle_aux(RaftState, Tag, Cmd, AuxV3, RaAux)
                   last_checkpoint = element(8, AuxV3)
                  },
     handle_aux(RaftState, Tag, Cmd, AuxV4, RaAux);
+handle_aux(RaftState, Tag, Cmd, AuxV4, RaAux)
+  when element(1, AuxV4) == aux_v4 ->
+    DecayMs = persistent_term:get(rabbit_fifo_ingress_decay_ms,
+                                  ?DEFAULT_INGRESS_DECAY_MS),
+    AuxV5 = #?AUX{name = element(2, AuxV4),
+                  last_decorators_state = element(3, AuxV4),
+                  last_consumer_timeout = element(4, AuxV4),
+                  gc = element(5, AuxV4),
+                  tick_pid = element(6, AuxV4),
+                  cache = element(7, AuxV4),
+                  last_checkpoint = element(8, AuxV4),
+                  ingress = #ingress_aux{decay_ms = DecayMs}},
+    handle_aux(RaftState, Tag, Cmd, AuxV5, RaAux);
 handle_aux(leader, cast, eval,
            #?AUX{last_decorators_state = LastDec,
                  last_consumer_timeout = LastConTimeout0,
@@ -1348,22 +1403,25 @@ handle_aux(_RaftState, cast, {#return{msg_ids = MsgIds,
             %% for returns with a delivery limit set we can just return as before
             {no_reply, Aux0, RaAux0, [{append, Ret, {notify, Corr, Pid}}]}
     end;
-handle_aux(leader, _, {handle_tick, [QName, Overview0, Nodes]},
-           #?AUX{tick_pid = Pid} = Aux, RaAux) ->
-    Overview = Overview0#{members_info => ra_aux:members_info(RaAux)},
-    NewPid =
-        case process_is_alive(Pid) of
-            false ->
-                %% No active TICK pid
-                %% this function spawns and returns the tick process pid
-                rabbit_quorum_queue:handle_tick(QName, Overview, Nodes);
-            true ->
-                %% Active TICK pid, do nothing
-                Pid
-        end,
 
-    %% TODO: check consumer timeouts
-    {no_reply, Aux#?AUX{tick_pid = NewPid}, RaAux, []};
+handle_aux(RaftState, _, {handle_tick, [QName, Overview0, Nodes]},
+           #?AUX{tick_pid = Pid, ingress = Ingress0} = Aux, RaAux) ->
+    Overview = Overview0#{members_info => ra_aux:members_info(RaAux)},
+    Ingress = update_ingress(Overview0, Nodes, Ingress0),
+    Aux1 = Aux#?AUX{ingress = Ingress},
+    case RaftState of
+        leader ->
+            NewPid =
+                case process_is_alive(Pid) of
+                    false ->
+                        rabbit_quorum_queue:handle_tick(QName, Overview, Nodes);
+                    true ->
+                        Pid
+                end,
+            {no_reply, Aux1#?AUX{tick_pid = NewPid}, RaAux, []};
+        _ ->
+            {no_reply, Aux1, RaAux, []}
+    end;
 handle_aux(_, _, {get_checked_out, ConsumerKey, MsgIds}, Aux0, RaAux0) ->
     #?STATE{cfg = #cfg{},
             consumers = Consumers} = ra_aux:machine_state(RaAux0),
@@ -1460,6 +1518,10 @@ handle_aux(leader, _, {dlx, setup}, Aux, RaAux) ->
 handle_aux(_, _, {dlx, teardown, Pid}, Aux, RaAux) ->
     terminate_dlx_worker(Pid),
     {no_reply, Aux, RaAux};
+handle_aux(_, {call, _From}, get_ingress_rates,
+           #?AUX{ingress = Ingress} = Aux, RaAux) ->
+    Rates = compute_ingress_rates(Ingress),
+    {reply, {ok, Rates}, Aux, RaAux};
 handle_aux(_, _, Unhandled, Aux, RaAux) ->
     #?STATE{cfg = #cfg{resource = QR}} = ra_aux:machine_state(RaAux),
     ?LOG_DEBUG("~ts: rabbit_fifo: unhandled aux command ~P",
@@ -1907,12 +1969,24 @@ maybe_return_all(#{system_time := Ts} = Meta, ConsumerKey,
              Effects}
     end.
 
+node_of(undefined) -> undefined;
+node_of(Pid) when is_pid(Pid) -> node(Pid).
+
+bump_ingress(Node, Size, Map) ->
+    maps:update_with(Node, fun(V) -> V + Size end, Size, Map).
+
 apply_enqueue(#{index := RaftIdx,
                 system_time := Ts} = Meta, From,
               Seq, RawMsg, Size, State0) ->
     case maybe_enqueue(RaftIdx, Ts, From, Seq, RawMsg, Size, [], State0) of
         {ok, State1, Effects1} ->
-            checkout(Meta, State0, State1, Effects1);
+            {MetaSize, BodySize} = Size,
+            TotalSize = MetaSize + BodySize,
+            IngressByNode = State1#?STATE.ingress_bytes_by_node,
+            State2 = State1#?STATE{
+                ingress_bytes_by_node =
+                    bump_ingress(node_of(From), TotalSize, IngressByNode)},
+            checkout(Meta, State0, State2, Effects1);
         {out_of_sequence, State, Effects} ->
             {State, not_enqueued, Effects};
         {duplicate, State, Effects} ->
diff --git a/deps/rabbit/src/rabbit_fifo.hrl b/deps/rabbit/src/rabbit_fifo.hrl
@@ -295,7 +295,12 @@
          last_active :: option(non_neg_integer()),
          msg_cache :: option({ra:index(), raw_msg()}),
          %% delayed retry messages awaiting redelivery
-         delayed = #delayed{} :: #delayed{}
+         delayed = #delayed{} :: #delayed{},
+         %% bytes enqueued, accumulated per originating node since the
+         %% creation of this state machine. Monotonically non-decreasing
+         %% on every node (under leader replication). Used by Tanzu's
+         %% leader-migration plugin to estimate per-node ingress.
+         ingress_bytes_by_node = #{} :: #{node() | undefined => non_neg_integer()}
         }).
 
 -type config() :: #{name := atom(),
diff --git a/deps/rabbit/test/quorum_queue_SUITE.erl b/deps/rabbit/test/quorum_queue_SUITE.erl
@@ -4670,7 +4670,9 @@ purge(Config) ->
 
     {'queue.purge_ok', 2} = amqp_channel:call(Ch, #'queue.purge'{queue = QQ}),
 
-    ?assertEqual([0], dirty_query([Server], RaName, fun rabbit_fifo:query_messages_total/1)).
+    ?assertMatch(#{num_messages := 0},  machine_overview({RaName, Server})),
+    ok.
+    % ?assertEqual([0], dirty_query([Server], RaName, fun rabbit_fifo:query_messages_total/1)).
 
 peek(Config) ->
     [Server | _] = rabbit_ct_broker_helpers:get_node_configs(Config, nodename),
diff --git a/deps/rabbit/test/rabbit_fifo_SUITE.erl b/deps/rabbit/test/rabbit_fifo_SUITE.erl
@@ -4030,7 +4030,7 @@ machine_version_test(Config) ->
                   consumers = #{3 := #consumer{cfg = #consumer_cfg{priority = 0}}},
                   service_queue = S,
                   messages = Msgs}, ok,
-     [_|_]} = apply(meta(Config, Idx), {machine_version, 7, 8}, S1),
+     [_|_]} = apply(meta(Config, Idx), {machine_version, 7, 9}, S1),
 
     ?assertEqual(1, rabbit_fifo_pq:len(Msgs)),
     ?assert(priority_queue:is_queue(S)),
@@ -4058,16 +4058,16 @@ machine_version_waiting_consumer_test(Config) ->
                                                  #consumer_cfg{priority = 0}}},
                   service_queue = S,
                   messages = Msgs}, ok, _} = apply(meta(Config, Idx),
-                                                    {machine_version, 7, 8}, S1),
+                                                    {machine_version, 7, 9}, S1),
     %% validate message conversion to lqueue
     ?assertEqual(0, rabbit_fifo_pq:len(Msgs)),
     ?assert(priority_queue:is_queue(S)),
     ?assertEqual(1, priority_queue:len(S)),
     ok.
 
-convert_v7_to_v8_test(Config) ->
+convert_v7_to_v9_test(Config) ->
     ConfigV7 = [{machine_version, 7} | Config],
-    ConfigV8 = [{machine_version, 8} | Config],
+    ConfigV9 = [{machine_version, 9} | Config],
 
     EPid = test_util:fake_pid(node()),
     Pid1 = test_util:fake_pid(node()),
@@ -4092,7 +4092,7 @@ convert_v7_to_v8_test(Config) ->
     {StateV7, _} = run_log(rabbit_fifo_v7, ConfigV7, Init, Entries,
                            fun (_) -> true end),
     {#rabbit_fifo{consumers = Consumers}, ok, _} =
-        apply(meta(ConfigV8, ?LINE), {machine_version, 7, 8}, StateV7),
+        apply(meta(ConfigV9, ?LINE), {machine_version, 7, 9}, StateV7),
 
     ?assertMatch(#consumer{status = {suspected_down, up}},
                  maps:get(Cid1, Consumers)),
@@ -4927,6 +4927,36 @@ query_single_active_consumer_consumer_info_test(Config) ->
     ok.
 
 
+%% Ingress tracking tests
+
+ingress_bytes_by_node_accumulates_on_enqueue_test(Config) ->
+    S0 = test_init(ingress_accumulates),
+    Msg = mk_mc(<<"hello">>),
+    Pid = self(),
+    Enq = make_enqueue(Pid, 1, Msg),
+    {S1, _, _} = apply(meta(Config, 1), Enq, S0),
+    Ingress = S1#rabbit_fifo.ingress_bytes_by_node,
+    ?assert(maps:size(Ingress) > 0),
+    ok.
+
+ingress_bytes_by_node_survives_snapshot_test(Config) ->
+    S0 = test_init(ingress_snapshot),
+    Msg = mk_mc(<<"test">>),
+    Pid = self(),
+    Enq = make_enqueue(Pid, 1, Msg),
+    {S1, _, _} = apply(meta(Config, 1), Enq, S0),
+    Ingress = S1#rabbit_fifo.ingress_bytes_by_node,
+    ?assert(maps:size(Ingress) > 0),
+    %% Simulate snapshot/restore
+    S2 = S1,
+    Ingress2 = S2#rabbit_fifo.ingress_bytes_by_node,
+    ?assertEqual(Ingress, Ingress2),
+    ok.
+
+
+%% Ingress tracking tests end
+
+
 %% Utility
 
 init(Conf) -> rabbit_fifo:init(Conf).
diff --git a/deps/rabbit/test/rabbit_fifo_prop_SUITE.erl b/deps/rabbit/test/rabbit_fifo_prop_SUITE.erl
@@ -10,7 +10,7 @@
 -include_lib("rabbit_common/include/rabbit_framing.hrl").
 -include_lib("rabbit_common/include/rabbit.hrl").
 
--define(MACHINE_VERSION, 8).
+-define(MACHINE_VERSION, 9).
 
 %%%===================================================================
 %%% Common Test callbacks
@@ -84,7 +84,8 @@ all_tests() ->
      dlx_09,
      single_active_ordering_02,
      two_nodes_same_otp_version,
-     two_nodes_different_otp_version
+     two_nodes_different_otp_version,
+     ingress_bytes_by_node_accumulation
     ].
 
 groups() ->
@@ -1126,6 +1127,29 @@ is_same_otp_version(ConfigOrNode) ->
     ct:pal("Our CT node runs OTP ~s, other node runs OTP ~s", [OurOTP, OtherOTP]),
     OurOTP =:= OtherOTP.
 
+ingress_bytes_by_node_accumulation(_Config) ->
+    Size = 500,
+    run_proper(
+      fun () ->
+              InitConf = config(?FUNCTION_NAME, undefined, undefined, false, undefined),
+              ?FORALL(O, ?LET(Ops, log_gen_different_nodes(Size), expand(Ops, InitConf)),
+                      begin
+                          Indexes = lists:seq(1, length(O)),
+                          Entries = lists:zip(Indexes, O),
+                          InitState = test_init(InitConf),
+                          {State1, _Effs1} = run_log(InitState, Entries),
+                          IngressByNode = State1#rabbit_fifo.ingress_bytes_by_node,
+                          %% Verify the map is properly populated
+                          IsMap = is_map(IngressByNode),
+                          %% Verify all values are non-negative
+                          ValidValues = maps:fold(
+                            fun(_, V, Acc) ->
+                                    Acc andalso is_integer(V) andalso V >= 0
+                            end, true, IngressByNode),
+                          IsMap andalso ValidValues
+                      end)
+      end, [], Size).
+
 two_nodes(Node) ->
     Size = 100,
     run_proper(
@@ -1511,7 +1535,7 @@ different_nodes_prop(Node, Conf, Commands) ->
     Entries = lists:zip(Indexes, Commands),
     InitState = test_init(Conf),
     Fun = fun(_) -> true end,
-    MachineVersion = 8,
+    MachineVersion = rabbit_fifo:version(),
 
     {State1, _Effs1} = run_log(InitState, Entries, Fun, MachineVersion),
     {State2, _Effs2} = erpc:call(Node, ?MODULE, run_log,
@@ -1598,8 +1622,8 @@ valid_simple_prefetch(_, _, _, _, _) ->
     true.
 
 upgrade_prop(Conf, Commands) ->
-    FromVersion = 7,
-    ToVersion = 8,
+    FromVersion = 8,
+    ToVersion = 9,
     FromMod = rabbit_fifo:which_module(FromVersion),
     ToMod = rabbit_fifo:which_module(ToVersion),
     Indexes = lists:seq(1, length(Commands)),
