Report message bytes in quorum queue stats

[#161505138]

Author: dcorbacho

src/rabbit_fifo.erl

@@ -41,13 +41,16 @@
          query_consumer_count/1,
          usage/1,
 
+         zero/1,
+
          %% misc
          dehydrate_state/1
         ]).
 
 -ifdef(TEST).
 -export([
-         metrics_handler/1
+         metrics_handler/1,
+         one/1
         ]).
 -endif.
 
@@ -186,6 +189,7 @@
          cancel_consumer_handler :: maybe(applied_mfa()),
          become_leader_handler :: maybe(applied_mfa()),
          metrics_handler :: maybe(applied_mfa()),
+         message_size_handler :: maybe(applied_mfa()),
          %% This is a special field that is only used for snapshots
          %% It represents the number of queued messages at the time the
          %% dehydrated snapshot state was cached.
@@ -196,7 +200,9 @@
          %% it instead takes messages from the `messages' map.
          %% This is done so that consumers are still served in a deterministic
          %% order on recovery.
-         prefix_msg_count = 0 :: non_neg_integer()
+         prefix_msg_count = 0 :: non_neg_integer(),
+         msg_bytes_enqueue = 0 :: non_neg_integer(),
+         msg_bytes_checkout = 0 :: non_neg_integer()
         }).
 
 -opaque state() :: #state{}.
@@ -206,7 +212,8 @@
                     become_leader_handler => applied_mfa(),
                     cancel_consumer_handler => applied_mfa(),
                     metrics_handler => applied_mfa(),
-                    shadow_copy_interval => non_neg_integer()}.
+                    shadow_copy_interval => non_neg_integer(),
+                    message_size_handler => applied_mfa()}.
 
 -export_type([protocol/0,
               delivery/0,
@@ -231,11 +238,16 @@ update_state(Conf, State) ->
     BLH = maps:get(become_leader_handler, Conf, undefined),
     MH = maps:get(metrics_handler, Conf, undefined),
     SHI = maps:get(shadow_copy_interval, Conf, ?SHADOW_COPY_INTERVAL),
+    MSH = maps:get(message_size_handler, Conf, {?MODULE, zero, []}),
     State#state{dead_letter_handler = DLH,
                 cancel_consumer_handler = CCH,
                 become_leader_handler = BLH,
                 metrics_handler = MH,
-                shadow_copy_interval = SHI}.
+                shadow_copy_interval = SHI,
+                message_size_handler = MSH}.
+
+zero(_) ->
+    0.
 
 % msg_ids are scoped per consumer
 % ra_indexes holds all raft indexes for enqueues currently on queue
@@ -245,7 +257,8 @@ update_state(Conf, State) ->
 apply(#{index := RaftIdx}, {enqueue, From, Seq, RawMsg}, Effects0, State00) ->
     case maybe_enqueue(RaftIdx, From, Seq, RawMsg, Effects0, State00) of
         {ok, State0, Effects1} ->
-            {State, Effects, ok} = checkout(State0, Effects1),
+            {State, Effects, ok} = checkout(add_bytes_enqueue(RawMsg, State0),
+                                            Effects1),
             {append_to_master_index(RaftIdx, State), Effects, ok};
         {duplicate, State, Effects} ->
             {State, Effects, ok}
@@ -506,12 +519,16 @@ tick(_Ts, #state{name = Name,
                  messages = Messages,
                  ra_indexes = Indexes,
                  metrics_handler = MH,
-                 consumers = Cons} = State) ->
+                 consumers = Cons,
+                 msg_bytes_enqueue = EnqueueBytes,
+                 msg_bytes_checkout = CheckoutBytes} = State) ->
     Metrics = {Name,
                maps:size(Messages), % Ready
                num_checked_out(State), % checked out
                rabbit_fifo_index:size(Indexes), %% Total
-               maps:size(Cons)}, % Consumers
+               maps:size(Cons), % Consumers
+               EnqueueBytes,
+               CheckoutBytes},
     case MH of
         undefined ->
             [{aux, emit}];
@@ -523,13 +540,18 @@ tick(_Ts, #state{name = Name,
 overview(#state{consumers = Cons,
                 enqueuers = Enqs,
                 messages = Messages,
-                ra_indexes = Indexes} = State) ->
+                ra_indexes = Indexes,
+                msg_bytes_enqueue = EnqueueBytes,
+                msg_bytes_checkout = CheckoutBytes
+               } = State) ->
     #{type => ?MODULE,
       num_consumers => maps:size(Cons),
       num_checked_out => num_checked_out(State),
       num_enqueuers => maps:size(Enqs),
       num_ready_messages => maps:size(Messages),
-      num_messages => rabbit_fifo_index:size(Indexes)}.
+      num_messages => rabbit_fifo_index:size(Indexes),
+      enqueue_message_bytes => EnqueueBytes,
+      checkout_message_bytes => CheckoutBytes}.
 
 -spec get_checked_out(consumer_id(), msg_id(), msg_id(), state()) ->
     [delivery_msg()].
@@ -766,12 +788,17 @@ complete_and_checkout(IncomingRaftIdx, MsgIds, ConsumerId,
     Checked = maps:without(MsgIds, Checked0),
     Discarded = maps:with(MsgIds, Checked0),
     MsgRaftIdxs = [RIdx || {_, {RIdx, _}} <- maps:values(Discarded)],
+    State1 = lists:foldl(fun({_, {_, {_, RawMsg}}}, Acc) ->
+                                 add_bytes_settle(RawMsg, Acc);
+                            (_, Acc) ->
+                                 Acc
+                         end, State0, maps:values(Discarded)),
     %% need to pass the length of discarded as $prefix_msgs would be filtered
     %% by the above list comprehension
-    {State1, Effects1, _} = complete(ConsumerId, MsgRaftIdxs,
+    {State2, Effects1, _} = complete(ConsumerId, MsgRaftIdxs,
                                      maps:size(Discarded),
-                                     Con0, Checked, Effects0, State0),
-    {State, Effects, _} = checkout(State1, Effects1),
+                                     Con0, Checked, Effects0, State1),
+    {State, Effects, _} = checkout(State2, Effects1),
     % settle metrics are incremented separately
     update_smallest_raft_index(IncomingRaftIdx, Indexes0, State, Effects).
 
@@ -834,8 +861,9 @@ return_one(MsgNum, {RaftId, {Header0, RawMsg}},
                               1, Header0),
     Msg = {RaftId, {Header, RawMsg}},
     % this should not affect the release cursor in any way
-    State0#state{messages = maps:put(MsgNum, Msg, Messages),
-                 returns = queue:in(MsgNum, Returns)}.
+    add_bytes_return(RawMsg,
+                     State0#state{messages = maps:put(MsgNum, Msg, Messages),
+                                  returns = queue:in(MsgNum, Returns)}).
 
 return_all(State, Checked) ->
     maps:fold(fun (_, '$prefix_msg',
@@ -941,14 +969,17 @@ checkout_one(#state{service_queue = SQ0,
                             {Cons, SQ, []} = % we expect no effects
                                 update_or_remove_sub(ConsumerId, Con,
                                                      Cons0, SQ1, []),
-                            State = State0#state{service_queue = SQ,
-                                                 messages = Messages,
-                                                 prefix_msg_count = PrefMsgC,
-                                                 consumers = Cons},
-                            Msg = case ConsumerMsg of
-                                      '$prefix_msg' -> '$prefix_msg';
-                                      {_, {_, M}} -> M
-                                  end,
+                            State1 = State0#state{service_queue = SQ,
+                                                  messages = Messages,
+                                                  prefix_msg_count = PrefMsgC,
+                                                  consumers = Cons},
+                            {State, Msg} =
+                                case ConsumerMsg of
+                                    '$prefix_msg' ->
+                                        {State1, '$prefix_msg'};
+                                    {_, {_, {_, RawMsg} = M}} ->
+                                        {add_bytes_checkout(RawMsg, State1), M}
+                                end,
                             {success, ConsumerId, Next, Msg, State};
                         error ->
                             %% consumer did not exist but was queued, recurse
@@ -1052,6 +1083,29 @@ dehydrate_consumer(#consumer{checked_out = Checked0} = Con) ->
     Checked = maps:map(fun (_, _) -> '$prefix_msg' end, Checked0),
     Con#consumer{checked_out = Checked}.
 
+add_bytes_enqueue(Msg, #state{message_size_handler = {M, F, A},
+                              msg_bytes_enqueue = Enqueue} = State) ->
+    Bytes = apply(M, F, [Msg | A]),
+    State#state{msg_bytes_enqueue = Enqueue + Bytes}.
+
+add_bytes_checkout(Msg, #state{message_size_handler = {M, F, A},
+                               msg_bytes_checkout = Checkout,
+                               msg_bytes_enqueue = Enqueue } = State) ->
+    Bytes = apply(M, F, [Msg | A]),
+    State#state{msg_bytes_checkout = Checkout + Bytes,
+                msg_bytes_enqueue = Enqueue - Bytes}.
+
+add_bytes_settle(Msg, #state{message_size_handler = {M, F, A},
+                             msg_bytes_checkout = Checkout} = State) ->
+    Bytes = apply(M, F, [Msg | A]),
+    State#state{msg_bytes_checkout = Checkout - Bytes}.
+
+add_bytes_return(Msg, #state{message_size_handler = {M, F, A},
+                             msg_bytes_checkout = Checkout,
+                             msg_bytes_enqueue = Enqueue} = State) ->
+    Bytes = apply(M, F, [Msg | A]),
+    State#state{msg_bytes_checkout = Checkout - Bytes,
+                msg_bytes_enqueue = Enqueue + Bytes}.
 
 -ifdef(TEST).
 -include_lib("eunit/include/eunit.hrl").
@@ -1079,6 +1133,12 @@ test_init(Name) ->
            shadow_copy_interval => 0,
            metrics_handler => {?MODULE, metrics_handler, []}}).
 
+test_init(Name, MsgSizeHandler) ->
+    init(#{name => Name,
+           shadow_copy_interval => 0,
+           metrics_handler => {?MODULE, metrics_handler, []},
+           message_size_handler => MsgSizeHandler}).
+
 metrics_handler(_) ->
     ok.
 
@@ -1416,16 +1476,19 @@ discarded_message_with_dead_letter_handler_emits_mod_call_effect_test() ->
                 Effects2),
     ok.
 
+one(_) ->
+    1.
+
 tick_test() ->
     Cid = {<<"c">>, self()},
     Cid2 = {<<"c2">>, self()},
-    {S0, _} = enq(1, 1, fst, test_init(test)),
+    {S0, _} = enq(1, 1, fst, test_init(test, {?MODULE, one, []})),
     {S1, _} = enq(2, 2, snd, S0),
     {S2, {MsgId, _}} = deq(3, Cid, unsettled, S1),
     {S3, {_, _}} = deq(4, Cid2, unsettled, S2),
     {S4, _, _} = apply(meta(5), {return, [MsgId], Cid}, [], S3),
 
-    [{mod_call, _, _, [{test, 1, 1, 2, 1}]}, {aux, emit}] = tick(1, S4),
+    [{mod_call, _, _, [{test, 1, 1, 2, 1, 1, 1}]}, {aux, emit}] = tick(1, S4),
     ok.
 
 enq_deq_snapshot_recover_test() ->

src/rabbit_quorum_queue.erl

@@ -36,6 +36,7 @@
 -export([requeue/3]).
 -export([policy_changed/2]).
 -export([cleanup_data_dir/0]).
+-export([message_size/1]).
 
 -include_lib("rabbit_common/include/rabbit.hrl").
 -include_lib("stdlib/include/qlc.hrl").
@@ -159,7 +160,12 @@ ra_machine_config(Q = #amqqueue{name = QName}) ->
     #{dead_letter_handler => dlx_mfa(Q),
       cancel_consumer_handler => {?MODULE, cancel_consumer, [QName]},
       become_leader_handler => {?MODULE, become_leader, [QName]},
-      metrics_handler => {?MODULE, update_metrics, [QName]}}.
+      metrics_handler => {?MODULE, update_metrics, [QName]},
+      message_size_handler => {?MODULE, message_size, []}}.
+
+message_size(#basic_message{content = Content} = Msg) ->
+    #content{payload_fragments_rev = PFR} = Content,
+    iolist_size(PFR).
 
 cancel_consumer_handler(QName, {ConsumerTag, ChPid}, _Name) ->
     Node = node(ChPid),
@@ -207,14 +213,17 @@ rpc_delete_metrics(QName) ->
     ets:delete(queue_metrics, QName),
     ok.
 
-update_metrics(QName, {Name, MR, MU, M, C}) ->
+update_metrics(QName, {Name, MR, MU, M, C, MsgBytesReady, MsgBytesUnack}) ->
     R = reductions(Name),
     rabbit_core_metrics:queue_stats(QName, MR, MU, M, R),
     Util = case C of
                0 -> 0;
                _ -> rabbit_fifo:usage(Name)
            end,
-    Infos = [{consumers, C}, {consumer_utilisation, Util} | infos(QName)],
+    Infos = [{consumers, C}, {consumer_utilisation, Util},
+             {message_bytes_ready, MsgBytesReady},
+             {message_bytes_unacknowledged, MsgBytesUnack},
+             {message_bytes, MsgBytesReady + MsgBytesUnack} | infos(QName)],
     rabbit_core_metrics:queue_stats(QName, Infos),
     rabbit_event:notify(queue_stats, Infos ++ [{name, QName},
                                                {messages, M},

test/quorum_queue_SUITE.erl

@@ -119,7 +119,8 @@ all_tests() ->
      delete_immediately,
      delete_immediately_by_resource,
      consume_redelivery_count,
-     subscribe_redelivery_count
+     subscribe_redelivery_count,
+     message_bytes_metrics
     ].
 
 %% -------------------------------------------------------------------
@@ -2044,6 +2045,65 @@ consume_redelivery_count(Config) ->
                                         multiple     = false,
                                         requeue      = true}).
 
+message_bytes_metrics(Config) ->
+    [Server | _] = Servers = rabbit_ct_broker_helpers:get_node_configs(Config, nodename),
+
+    Ch = rabbit_ct_client_helpers:open_channel(Config, Server),
+    QQ = ?config(queue_name, Config),
+    ?assertEqual({'queue.declare_ok', QQ, 0, 0},
+                 declare(Ch, QQ, [{<<"x-queue-type">>, longstr, <<"quorum">>}])),
+
+    RaName = ra_name(QQ),
+    {ok, _, {_, Leader}} = ra:members({RaName, Server}),
+    QRes = rabbit_misc:r(<<"/">>, queue, QQ),
+
+    publish(Ch, QQ),
+
+    wait_for_messages_ready(Servers, RaName, 1),
+    wait_for_messages_pending_ack(Servers, RaName, 0),
+    wait_until(fun() ->
+                       {3, 3, 0} == get_message_bytes(Leader, QRes)
+               end),
+
+    subscribe(Ch, QQ, false),
+
+    wait_for_messages_ready(Servers, RaName, 0),
+    wait_for_messages_pending_ack(Servers, RaName, 1),
+    wait_until(fun() ->
+                       {3, 0, 3} == get_message_bytes(Leader, QRes)
+               end),
+
+    receive
+        {#'basic.deliver'{delivery_tag = DeliveryTag,
+                          redelivered  = false}, _} ->
+            amqp_channel:cast(Ch, #'basic.nack'{delivery_tag = DeliveryTag,
+                                                multiple     = false,
+                                                requeue      = false}),
+            wait_for_messages_ready(Servers, RaName, 0),
+            wait_for_messages_pending_ack(Servers, RaName, 0),
+            wait_until(fun() ->
+                               {0, 0, 0} == get_message_bytes(Leader, QRes)
+                       end)
+    end,
+
+    %% Let's publish and then close the consumer channel. Messages must be
+    %% returned to the queue
+    publish(Ch, QQ),
+
+    wait_for_messages_ready(Servers, RaName, 0),
+    wait_for_messages_pending_ack(Servers, RaName, 1),
+    wait_until(fun() ->
+                       {3, 0, 3} == get_message_bytes(Leader, QRes)
+               end),
+
+    rabbit_ct_client_helpers:close_channel(Ch),
+
+    wait_for_messages_ready(Servers, RaName, 1),
+    wait_for_messages_pending_ack(Servers, RaName, 0),
+    wait_until(fun() ->
+                       {3, 3, 0} == get_message_bytes(Leader, QRes)
+               end).
+
 %%----------------------------------------------------------------------------
 
 declare(Ch, Q) ->
@@ -2221,3 +2281,13 @@ delete_queues() ->
 
 stop_node(Config, Server) ->
     rabbit_ct_broker_helpers:rabbitmqctl(Config, Server, ["stop"]).
+
+get_message_bytes(Leader, QRes) ->
+    case rpc:call(Leader, ets, lookup, [queue_metrics, QRes]) of
+        [{QRes, Props, _}] ->
+            {proplists:get_value(message_bytes, Props),
+             proplists:get_value(message_bytes_ready, Props),
+             proplists:get_value(message_bytes_unacknowledged, Props)};
+        _ ->
+            []
+    end.