Add per-node ingress byte tracking to rabbit_fifo v9

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.

Author: kjnilsson

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} ->

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(),

deps/rabbit/test/quorum_queue_SUITE.erl

@@ -4670,7 +4670,8 @@ 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.
 
 peek(Config) ->
     [Server | _] = rabbit_ct_broker_helpers:get_node_configs(Config, nodename),

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).

deps/rabbit/test/rabbit_fifo_dlx_integration_SUITE.erl

@@ -655,8 +655,8 @@ reject_publish_max_length_target_quorum_queue(Config) ->
                      amqp_channel:call(Ch, #'basic.get'{queue = TargetQ}),
                      30000)
      end || N <- lists:seq(1,4)],
-    eventually(?_assertMatch([{0, _}],
-                             dirty_query([Server], RaName, fun rabbit_fifo:query_stat_dlx/1)), 500, 10),
+    eventually(?_assertMatch(#{num_discarded := 0},
+                             machine_overview({RaName, Server}))),
     ?assertEqual(4, counted(messages_dead_lettered_expired_total, Config)),
     eventually(?_assertEqual(4, counted(messages_dead_lettered_confirmed_total, Config))).
 
@@ -706,8 +706,8 @@ reject_publish_down_target_quorum_queue(Config) ->
                          sets:add_element(Msg, S)
                  end, sets:new([{version, 2}]), lists:seq(1, 50)),
     ?assertEqual(50, sets:size(Received)),
-    eventually(?_assertMatch([{0, _}],
-                             dirty_query([Server], RaName, fun rabbit_fifo:query_stat_dlx/1)), 500, 10),
+    eventually(?_assertMatch(#{num_discarded := 0},
+                             machine_overview({RaName, Server}))),
     ?assertEqual(50, counted(messages_dead_lettered_expired_total, Config)),
     eventually(?_assertEqual(50, counted(messages_dead_lettered_confirmed_total, Config))).
 
@@ -999,3 +999,11 @@ counted(Metric, Config) ->
 metric(Metric, Counters) ->
     Metrics = maps:get(#{queue_type => rabbit_quorum_queue, dead_letter_strategy => at_least_once}, Counters),
     maps:get(Metric, Metrics).
+
+machine_overview(ServerId) when is_tuple(ServerId) ->
+    case ra:member_overview(ServerId) of
+        {ok, #{machine := Mac}, _} ->
+            Mac;
+        Err ->
+            Err
+    end.