QQ: record lengths of each priority queue

kjnilsson · kjnilsson · commit 235bd37ebaa0 · 2025-11-24T14:23:08.000Z
And emit the number of messages for each priority in
the overview map.
diff --git a/deps/rabbit/src/rabbit_fifo.erl b/deps/rabbit/src/rabbit_fifo.erl
@@ -449,7 +449,7 @@ apply_(#{index := Idx} = Meta,
                {Life, _Credit, credited} ->
                    {Life, credited}
            end,
-    Priority = get_priority(ConsumerMeta),
+    Priority = get_consumer_priority(ConsumerMeta),
     ConsumerKey = case consumer_key_from_id(ConsumerId, State0) of
                       {ok, K} ->
                           K;
@@ -1747,10 +1747,10 @@ maybe_enqueue(RaftIdx, Ts, undefined, undefined, RawMsg,
     Header0 = maybe_set_msg_ttl(RawMsg, Ts, Size, State0),
     Header = maybe_set_msg_delivery_count(RawMsg, Header0),
     Msg = make_msg(RaftIdx, Header),
-    PTag = priority_tag(RawMsg),
+    Priority = msg_priority(RawMsg),
     State = State0#?STATE{msg_bytes_enqueue = Enqueue + Size,
                           messages_total = Total + 1,
-                          messages = rabbit_fifo_pq:in(PTag, Msg, Messages)
+                          messages = rabbit_fifo_pq:in(Priority, Msg, Messages)
                          },
     {ok, State, Effects};
 maybe_enqueue(RaftIdx, Ts, From, MsgSeqNo, RawMsg,
@@ -1780,10 +1780,10 @@ maybe_enqueue(RaftIdx, Ts, From, MsgSeqNo, RawMsg,
                            false ->
                                undefined
                        end,
-            PTag = priority_tag(RawMsg),
+            Priority = msg_priority(RawMsg),
             State = State0#?STATE{msg_bytes_enqueue = BytesEnqueued + Size,
                                   messages_total = Total + 1,
-                                  messages = rabbit_fifo_pq:in(PTag, Msg, Messages),
+                                  messages = rabbit_fifo_pq:in(Priority, Msg, Messages),
                                   enqueuers = Enqueuers0#{From => Enq},
                                   msg_cache = MsgCache
                                  },
@@ -2835,17 +2835,17 @@ is_expired(Ts, #?STATE{cfg = #cfg{expires = Expires},
 is_expired(_Ts, _State) ->
     false.
 
-get_priority(#{priority := Priority}) ->
+get_consumer_priority(#{priority := Priority}) ->
     Priority;
-get_priority(#{args := Args}) ->
+get_consumer_priority(#{args := Args}) ->
     %% fallback, v3 option
     case rabbit_misc:table_lookup(Args, <<"x-priority">>) of
         {_Type, Value} ->
             Value;
         _ ->
             0
     end;
-get_priority(_) ->
+get_consumer_priority(_) ->
     0.
 
 notify_decorators_effect(QName, MaxActivePriority, IsEmpty) ->
@@ -3037,21 +3037,23 @@ maps_search(Pred, {K, V, I}) ->
 maps_search(Pred, Map) when is_map(Map) ->
     maps_search(Pred, maps:next(maps:iterator(Map))).
 
-priority_tag(Msg) ->
+-define(DEFAULT_PRIORITY, 4).
+-define(MAX_PRIORITY, 31).
+
+msg_priority(Msg) ->
     case mc:is(Msg) of
         true ->
             case mc:priority(Msg) of
                 P when is_integer(P) ->
-                    min(P, 31);
+                    min(P, ?MAX_PRIORITY);
                 _ ->
-                    4
+                    ?DEFAULT_PRIORITY
             end;
         false ->
-            4
+            ?DEFAULT_PRIORITY
     end.
 
-do_snapshot(MacVer, Ts, Ch,
-               RaAux, DiscardedBytes, Force)
+do_snapshot(MacVer, Ts, Ch, RaAux, DiscardedBytes, Force)
   when element(1, Ch) == checkpoint andalso
        is_integer(MacVer)  andalso
        MacVer >= 8 ->
diff --git a/deps/rabbit/src/rabbit_fifo_pq.erl b/deps/rabbit/src/rabbit_fifo_pq.erl
@@ -8,6 +8,7 @@
 -module(rabbit_fifo_pq).
 
 -include("rabbit_fifo.hrl").
+
 -export([
          new/0,
          in/3,
@@ -22,68 +23,56 @@
          overview/1
         ]).
 
-% -define(NON_EMPTY, {_, [_|_]}).
--define(EMPTY, {[], []}).
+-define(STATE, pq).
+-define(EMPTY, {0, [], []}).
 
-%% a weighted priority queue with only two priorities
+%% supports 32 priorities, needs to be a power of 2 to support the De Bruijn
+%% lookup method. 64 would push the bitmap into an erlang big number so we
+%% have to settle for 32
 -type priority() :: 0..31.
--type queue() :: {list(msg()), list(msg())}.
+-type queue() :: {non_neg_integer(), list(msg()), list(msg())}.
 
--record(?MODULE, {buckets = #{} :: #{priority() => queue()},
-                  len = 0 :: non_neg_integer(),
-                  bitmap = 0 :: integer()}).
+-record(?STATE, {buckets = #{} :: #{priority() => queue()},
+                 len = 0 :: non_neg_integer(),
+                 bitmap = 0 :: integer()}).
 
--opaque state() :: #?MODULE{}.
+-opaque state() :: #?STATE{}.
 
 -export_type([state/0,
               priority/0]).
 
 -spec new() -> state().
 new() ->
-    #?MODULE{}.
+    #?STATE{}.
 
 -spec in(priority(), msg(), state()) -> state().
-in(Priority0, Item, #?MODULE{buckets = Buckets0,
-                             bitmap = Bitmap0,
-                             len = Len} = State)
+in(Priority0, Item, #?STATE{buckets = Buckets0,
+                            bitmap = Bitmap0,
+                            len = Len} = State)
   when Priority0 >= 0 andalso
        Priority0 =< 31 ->
     %% invert priority
     Priority = 31 - Priority0,
     case Buckets0 of
         #{Priority := Queue0} ->
             %% there are messages for the priority already
-            State#?MODULE{buckets = Buckets0#{Priority => in(Item, Queue0)},
-                          len = Len + 1};
+            State#?STATE{buckets = Buckets0#{Priority => in(Item, Queue0)},
+                         len = Len + 1};
         _ ->
             Bitmap = Bitmap0 bor (1 bsl Priority),
             %% there are no messages for the priority
-            State#?MODULE{buckets = Buckets0#{Priority => in(Item, ?EMPTY)},
-                          bitmap = Bitmap,
-                          len = Len + 1}
+            State#?STATE{buckets = Buckets0#{Priority => in(Item, ?EMPTY)},
+                         bitmap = Bitmap,
+                         len = Len + 1}
     end.
 
-first_set_bit(0) ->
-    32;
-first_set_bit(Bitmap) ->
-    count_trailing(Bitmap band -Bitmap).
-
--define(DEBRUIJN_SEQ, 16#077CB531).
--define(DEBRUIJN_LOOKUP,
-        {0, 1, 28, 2, 29, 14, 24, 3, 30, 22, 20, 15, 25, 17, 4, 8,
-         31, 27, 13, 23, 21, 19, 16, 7, 26, 12, 18, 6, 11, 5, 10, 9}).
-
-count_trailing(N) ->
-    Lookup = ((N * ?DEBRUIJN_SEQ) bsr 27) band 31,
-    element(Lookup + 1, ?DEBRUIJN_LOOKUP).
-
 -spec out(state()) ->
     empty | {msg(), state()}.
-out(#?MODULE{len = 0}) ->
+out(#?STATE{len = 0}) ->
     empty;
-out(#?MODULE{buckets = Buckets,
-             len = Len,
-             bitmap = Bitmap0} = State) ->
+out(#?STATE{buckets = Buckets,
+            len = Len,
+            bitmap = Bitmap0} = State) ->
     Priority = first_set_bit(Bitmap0),
     #{Priority := Q0} = Buckets,
     Msg = peek(Q0),
@@ -93,25 +82,25 @@ out(#?MODULE{buckets = Buckets,
             %% as we know the bit is set we just need to xor rather than
             %% create a mask then xor
             Bitmap = Bitmap0 bxor (1 bsl Priority),
-            {Msg, State#?MODULE{buckets = maps:remove(Priority, Buckets),
-                                len = Len - 1,
-                                bitmap = Bitmap}};
+            {Msg, State#?STATE{buckets = maps:remove(Priority, Buckets),
+                               len = Len - 1,
+                               bitmap = Bitmap}};
         Q ->
-            {Msg, State#?MODULE{buckets = maps:put(Priority, Q, Buckets),
-                                len = Len - 1}}
+            {Msg, State#?STATE{buckets = maps:put(Priority, Q, Buckets),
+                               len = Len - 1}}
     end.
 
 -spec get(state()) -> empty | msg().
-get(#?MODULE{len = 0}) ->
+get(#?STATE{len = 0}) ->
     empty;
-get(#?MODULE{buckets = Buckets,
-             bitmap = Bitmap}) ->
+get(#?STATE{buckets = Buckets,
+            bitmap = Bitmap}) ->
     Priority = first_set_bit(Bitmap),
     #{Priority := Q0} = Buckets,
     peek(Q0).
 
 -spec len(state()) -> non_neg_integer().
-len(#?MODULE{len = Len}) ->
+len(#?STATE{len = Len}) ->
     Len.
 
 -spec from_list([{priority(), term()}]) -> state().
@@ -131,17 +120,17 @@ from_lqueue(LQ) ->
                 end, new(), LQ).
 
 -spec indexes(state()) -> [ra:index()].
-indexes(#?MODULE{buckets = Buckets}) ->
+indexes(#?STATE{buckets = Buckets}) ->
     maps:fold(
-      fun (_P, {L1, L2}, Acc0) ->
+      fun (_P, {_, L1, L2}, Acc0) ->
               Acc = lists:foldl(fun msg_idx_fld/2, Acc0, L1),
               lists:foldl(fun msg_idx_fld/2, Acc, L2)
       end, [], Buckets).
 
 -spec get_lowest_index(state()) -> undefined | ra:index().
-get_lowest_index(#?MODULE{len = 0}) ->
+get_lowest_index(#?STATE{len = 0}) ->
     undefined;
-get_lowest_index(#?MODULE{buckets = Buckets}) ->
+get_lowest_index(#?STATE{buckets = Buckets}) ->
     lists:min(
       maps:fold(fun (_, Q, Acc) ->
                         case peek(Q) of
@@ -154,32 +143,38 @@ get_lowest_index(#?MODULE{buckets = Buckets}) ->
 
 -spec overview(state()) ->
     #{len := non_neg_integer(),
+      detail := #{priority() => pos_integer()},
       num_active_priorities := 0..32,
       lowest_index := ra:index()}.
-overview(#?MODULE{len = Len,
-                  buckets = Buckets} = State) ->
+overview(#?STATE{len = Len,
+                 buckets = Buckets} = State) ->
+    Detail  = maps:fold(fun (P0, {C, _, _}, Acc) ->
+                                P = 31-P0,
+                                Acc#{P => C}
+                        end, #{}, Buckets),
     #{len => Len,
+      detail => Detail,
       num_active_priorities => map_size(Buckets),
       lowest_index => get_lowest_index(State)}.
 
-%% internals
+%% INTERNAL
 
 %% invariant, if the queue is non empty so is the Out (right) list.
 in(X, ?EMPTY) ->
-    {[], [X]};
-in(X, {In, Out}) ->
-    {[X | In], Out}.
+    {1, [], [X]};
+in(X, {C, In, Out}) ->
+    {C+1, [X | In], Out}.
 
 peek(?EMPTY) ->
     empty;
-peek({_, [H | _]}) ->
+peek({_, _, [H | _]}) ->
     H.
 
-drop({In, [_]}) ->
+drop({C, In, [_]}) ->
     %% the last Out one
-    {[], lists:reverse(In)};
-drop({In, [_ | Out]}) ->
-    {In, Out}.
+    {C-1, [], lists:reverse(In)};
+drop({C, In, [_ | Out]}) ->
+    {C-1, In, Out}.
 
 msg_idx_fld(Msg, Acc) when is_list(Acc) ->
     [msg_idx(Msg) | Acc].
@@ -194,3 +189,17 @@ to_list(empty, Acc) ->
 to_list({Item, State}, Acc) ->
     to_list(out(State), [Item | Acc]).
 
+first_set_bit(0) ->
+    32;
+first_set_bit(Bitmap) ->
+    count_trailing(Bitmap band -Bitmap).
+
+-define(DEBRUIJN_SEQ, 16#077CB531).
+-define(DEBRUIJN_LOOKUP,
+        {0, 1, 28, 2, 29, 14, 24, 3, 30, 22, 20, 15, 25, 17, 4, 8,
+         31, 27, 13, 23, 21, 19, 16, 7, 26, 12, 18, 6, 11, 5, 10, 9}).
+
+count_trailing(N) ->
+    Lookup = ((N * ?DEBRUIJN_SEQ) bsr 27) band 31,
+    element(Lookup + 1, ?DEBRUIJN_LOOKUP).
+
diff --git a/deps/rabbit/test/rabbit_fifo_pq_SUITE.erl b/deps/rabbit/test/rabbit_fifo_pq_SUITE.erl
@@ -65,6 +65,11 @@ basics(_Config) ->
 
     ?assertEqual([1,2,3,4,5], lists:sort(rabbit_fifo_pq:indexes(Q1))),
     ?assertMatch(#{len := 5,
+                   detail := #{1 := 1,
+                               2 := 1,
+                               3 := 1,
+                               4 := 1,
+                               5 := 1},
                    num_active_priorities := 5,
                    lowest_index := 1}, rabbit_fifo_pq:overview(Q1)),
     {?MSG(5), Q2} = rabbit_fifo_pq:out(Q1),
@@ -112,14 +117,16 @@ get_lowest_index(_Config) ->
     ?assertEqual(1, rabbit_fifo_q:get_lowest_index(Q3)),
     ?assertEqual(2, rabbit_fifo_q:get_lowest_index(Q4)),
     ?assertEqual(3, rabbit_fifo_q:get_lowest_index(Q5)),
-    ?assertEqual(undefined, rabbit_fifo_q:get_lowest_index(Q6)).
+    ?assertEqual(undefined, rabbit_fifo_q:get_lowest_index(Q6)),
+    ok.
+
 
 property(_Config) ->
     run_proper(
       fun () ->
               ?FORALL(Ops, op_gen(256),
                       queue_prop(Ops))
-      end, [], 25),
+      end, [], 100),
     ok.
 
 queue_prop(Ops) ->
@@ -130,11 +137,15 @@ queue_prop(Ops) ->
 
     Sut0 = rabbit_fifo_pq:from_list(Ops),
     Out = rabbit_fifo_pq:to_list(Sut0),
+    #{detail := Detail,
+      len := Len} = rabbit_fifo_pq:overview(Sut0),
+    DetailSum = maps:fold(fun (_, C, Acc) -> Acc + C end, 0, Detail),
+    DetailSum == Len andalso
     [element(2, O) || O <- SortedOps] == Out.
 
 %%% helpers
 
--type item() :: {rabbit_fifo_pq:priority(), integer()}.
+-type item() :: {rabbit_fifo_pq:priority(), non_neg_integer()}.
 op_gen(Size) ->
     ?LET(Ops, resize(Size, list(item())), Ops).
 
