Local shovels: Add global counters

Author: dcorbacho

deps/rabbitmq_shovel/src/rabbit_local_shovel.erl

@@ -15,7 +15,15 @@
 -include_lib("rabbit/include/mc.hrl").
 -include("rabbit_shovel.hrl").
 
+-rabbit_boot_step({rabbit_global_local_shovel_counters,
+                   [{description, "global local shovel counters"},
+                    {mfa,         {?MODULE, boot_step,
+                                   []}},
+                    {requires,    rabbit_global_counters},
+                    {enables,     external_infrastructure}]}).
+
 -export([
+         boot_step/0,
          parse/2,
          connect_source/1,
          connect_dest/1,
@@ -53,12 +61,20 @@
 %% See rabbit_amqp_session.erl
 -define(INITIAL_DELIVERY_COUNT, 16#ff_ff_ff_ff - 4).
 -define(DEFAULT_MAX_LINK_CREDIT, 1000).
+-define(PROTOCOL, 'local-shovel').
 
 -record(pending_ack, {
                       delivery_tag,
                       msg_id
                      }).
 
+boot_step() ->
+    Labels = #{protocol => ?PROTOCOL},
+    rabbit_global_counters:init(Labels),
+    rabbit_global_counters:init(Labels#{queue_type => rabbit_classic_queue}),
+    rabbit_global_counters:init(Labels#{queue_type => rabbit_quorum_queue}),
+    rabbit_global_counters:init(Labels#{queue_type => rabbit_stream_queue}).
+
 parse(_Name, {source, Source}) ->
     Queue = parse_parameter(queue, fun parse_binary/1,
                             proplists:get_value(queue, Source)),
@@ -363,6 +379,7 @@ forward(Tag, Msg0, #{dest := #{current := #{queue_states := QState} = Current} =
     Msg = set_annotations(Msg0, Dest),
     RoutedQNames = route(Msg, Dest),
     Queues = rabbit_amqqueue:lookup_many(RoutedQNames),
+    messages_received(AckMode),
     case rabbit_queue_type:deliver(Queues, Msg, Options, QState) of
         {ok, QState1, Actions} ->
             State1 = State#{dest => Dest1#{current => Current1#{queue_states => QState1}}},
@@ -451,13 +468,15 @@ handle_queue_actions(Actions, State) ->
       end, State, Actions).
 
 handle_deliver(AckRequired, Msgs, State) when is_list(Msgs) ->
+    NumMsgs = length(Msgs),
     maybe_grant_credit(
       lists:foldl(
-        fun({_QName, _QPid, MsgId, _Redelivered, Mc}, S0) ->
+        fun({QName, _QPid, MsgId, _Redelivered, Mc}, #{source := #{current := #{queue_states := QStates }}} = S0) ->
+                messages_delivered(QName, QStates),
                 DeliveryTag = next_tag(S0),
                 S = record_pending(AckRequired, DeliveryTag, MsgId, increase_next_tag(S0)),
                 rabbit_shovel_behaviour:forward(DeliveryTag, Mc, S)
-        end, sent_delivery(State, length(Msgs)), Msgs)).
+        end, sent_delivery(State, NumMsgs), Msgs)).
 
 next_tag(#{source := #{current := #{next_tag := DeliveryTag}}}) ->
     DeliveryTag.
@@ -616,6 +635,7 @@ settle(Op, DeliveryTag, Multiple,
     {MsgIds, UAMQ} = collect_acks(UAMQ0, DeliveryTag, Multiple),
     case rabbit_queue_type:settle(QRef, Op, CTag, lists:reverse(MsgIds), QState0) of
         {ok, QState1, Actions} ->
+            messages_acknowledged(Op, QRef, QState1, MsgIds),
             State = State0#{source => Src#{current => Current#{queue_states => QState1,
                                                                unacked_message_q => UAMQ}}},
             handle_queue_actions(Actions, State);
@@ -739,12 +759,18 @@ handle_credit_reply({credit_reply, CTag, DeliveryCount, Credit, _Available, _Dra
                                    at_least_one_credit_req_in_flight => false}}
     end.
 
-process_routing_confirm(undefined, _, _, State) ->
+process_routing_confirm(undefined, _, [], State) ->
+    rabbit_global_counters:messages_unroutable_returned(?PROTOCOL, 1),
+    State;
+process_routing_confirm(undefined, _, QRefs, State) ->
+    rabbit_global_counters:messages_routed(?PROTOCOL, length(QRefs)),
     State;
 process_routing_confirm(MsgSeqNo, Tag, [], State)
   when is_integer(MsgSeqNo) ->
+    rabbit_global_counters:messages_unroutable_dropped(?PROTOCOL, 1),
     record_confirms([{MsgSeqNo, Tag}], State);
 process_routing_confirm(MsgSeqNo, Tag, QRefs, #{dest := Dst = #{unconfirmed := Unconfirmed}} = State) when is_integer(MsgSeqNo) ->
+    rabbit_global_counters:messages_routed(?PROTOCOL, length(QRefs)),
     State#{dest => Dst#{unconfirmed =>
                             rabbit_shovel_confirms:insert(MsgSeqNo, QRefs, Tag, Unconfirmed)}}.
 
@@ -781,8 +807,10 @@ send_nacks(Rs, Cs, State) ->
 send_confirms([], _, State) ->
     State;
 send_confirms([MsgSeqNo], _, State) ->
+    rabbit_global_counters:messages_confirmed(?PROTOCOL, 1),
     rabbit_shovel_behaviour:ack(MsgSeqNo, false, State);
 send_confirms(Cs, Rs, State) ->
+    rabbit_global_counters:messages_confirmed(?PROTOCOL, length(Cs)),
     coalesce_and_send(Cs, Rs,
                       fun(MsgSeqNo, Multiple, StateX) ->
                               rabbit_shovel_behaviour:ack(MsgSeqNo, Multiple, StateX)
@@ -833,3 +861,30 @@ decr_remaining(Num, State) ->
             _ = send_confirms_and_nacks(State),
             exit(R)
     end.
+
+messages_acknowledged(complete, QName, QS, MsgIds) ->
+    case rabbit_queue_type:module(QName, QS) of
+        {ok, QType} ->
+            rabbit_global_counters:messages_acknowledged(?PROTOCOL, QType, length(MsgIds));
+        _ ->
+            ok
+    end;
+messages_acknowledged(_, _, _, _) ->
+    ok.
+
+messages_received(AckMode) ->
+    rabbit_global_counters:messages_received(?PROTOCOL, 1),
+    case AckMode of
+        on_confirm ->
+            rabbit_global_counters:messages_received_confirm(?PROTOCOL, 1);
+        _ ->
+            ok
+    end.
+
+messages_delivered(QName, S0) ->
+    case rabbit_queue_type:module(QName, S0) of
+        {ok, QType} ->
+            rabbit_global_counters:messages_delivered(?PROTOCOL, QType, 1);
+        _ ->
+            ok
+    end.