AMQP 0-9-1: Handle per-queue-type disk alarms

This covers both network and direct connections for 0-9-1. We store a
set of the queue types which have been published into on both a channel
and connection level since blocking is done on the connection level but
only the channel knows what queue types have been published.

Then when the published queue types or the set of alarms changes, the
connection evaluates whether it is affected by the alarm. If not it may
publish but once a channel publishes to an alarmed queue type the
connection then blocks until the channel exits or the alarm clears.

Author: the-mikedavis

deps/amqp_client/src/amqp_gen_connection.erl

@@ -32,6 +32,7 @@
                 %% connection.block, connection.unblock handler
                 block_handler,
                 blocked_by = sets:new([{version, 2}]),
+                queue_types_published = sets:new([{version, 2}]),
                 closing = false %% #closing{} | false
                }).
 
@@ -214,22 +215,18 @@ handle_cast({register_blocked_handler, HandlerPid},
     {noreply, State1};
 handle_cast({conserve_resources, Source, Conserve},
             #state{blocked_by = BlockedBy} = State) ->
-    WasNotBlocked = sets:is_empty(BlockedBy),
     BlockedBy1 = case Conserve of
                      true ->
                          sets:add_element(Source, BlockedBy);
                      false ->
                          sets:del_element(Source, BlockedBy)
                  end,
     State1 = State#state{blocked_by = BlockedBy1},
-    case sets:is_empty(BlockedBy1) of
-        true ->
-            handle_method(#'connection.unblocked'{}, State1);
-        false when WasNotBlocked ->
-            handle_method(#'connection.blocked'{}, State1);
-        false ->
-            {noreply, State1}
-    end.
+    maybe_block(State, State1);
+handle_cast({channel_published_to_queue_type, _ChPid, QT},
+            #state{queue_types_published = QTs} = State) ->
+    State1 = State#state{queue_types_published = sets:add_element(QT, QTs)},
+    maybe_block(State, State1).
 
 %% @private
 handle_info({'DOWN', _, process, BlockHandler, Reason},
@@ -274,6 +271,24 @@ i(Item, #state{module = Mod, module_state = MState}) -> Mod:i(Item, MState).
 register_blocked_handler(Pid, HandlerPid) ->
     gen_server:cast(Pid, {register_blocked_handler, HandlerPid}).
 
+maybe_block(State0, State1) ->
+    WasBlocked = should_block(State0),
+    case should_block(State1) of
+        true when not WasBlocked ->
+            handle_method(#'connection.blocked'{}, State1);
+        false when WasBlocked ->
+            handle_method(#'connection.unblocked'{}, State1);
+        _ ->
+            {noreply, State1}
+    end.
+
+should_block(#state{blocked_by = BlockedBy, queue_types_published = QTs}) ->
+    lists:any(fun ({disk, QT}) ->
+                      sets:is_element(QT, QTs);
+                  (_Resource) ->
+                      true
+              end, sets:to_list(BlockedBy)).
+
 %%---------------------------------------------------------------------------
 %% Command handling
 %%---------------------------------------------------------------------------

deps/rabbit/src/rabbit_amqqueue.erl

@@ -67,7 +67,7 @@
 -export([deactivate_limit_all/2]).
 
 -export([prepend_extra_bcc/1]).
--export([queue/1, queue_names/1]).
+-export([queue/1, queue_names/1, queue_types/1]).
 
 -export([kill_queue/2, kill_queue/3, kill_queue_hard/2, kill_queue_hard/3]).
 -export([delete_transient_queues_on_node/1]).
@@ -2097,6 +2097,15 @@ queue_names(Queues) ->
                       amqqueue:get_name(Q)
               end, Queues).
 
+-spec queue_types([Q | {Q, route_infos()}]) ->
+    [rabbit_queue_type:queue_type()] when Q :: amqqueue:target().
+queue_types(Queues) ->
+    lists:map(fun({Q, RouteInfos}) when is_map(RouteInfos) ->
+                      amqqueue:get_type(Q);
+                 (Q) ->
+                      amqqueue:get_type(Q)
+              end, Queues).
+
 -spec lookup_extra_bcc(amqqueue:target(), binary()) ->
     [amqqueue:target()].
 lookup_extra_bcc(Q, BCCName) ->

deps/rabbit/src/rabbit_channel.erl

@@ -169,6 +169,8 @@
              delivery_flow :: flow | noflow,
              interceptor_state,
              queue_states,
+             queue_types_published = sets:new([{version, 2}]) ::
+                                     sets:set(rabbit_queue_type:queue_type()),
              tick_timer,
              publishing_mode = false :: boolean()
             }).
@@ -2092,14 +2094,16 @@ deliver_to_queues(XName,
         {ok, QueueStates, Actions} ->
             rabbit_global_counters:messages_routed(amqp091, length(Qs)),
             QueueNames = rabbit_amqqueue:queue_names(Qs),
+            QueueTypes = rabbit_amqqueue:queue_types(Qs),
             %% NB: the order here is important since basic.returns must be
             %% sent before confirms.
             ok = process_routing_mandatory(Mandatory, RoutedToQueues, Message, XName, State0),
             MsgSeqNo = maps:get(correlation, Options, undefined),
             State1 = process_routing_confirm(MsgSeqNo, QueueNames, XName, State0),
+            State2 = notify_published_queue_types(QueueTypes, State1),
             %% Actions must be processed after registering confirms as actions may
             %% contain rejections of publishes
-            State = handle_queue_actions(Actions, State1#ch{queue_states = QueueStates}),
+            State = handle_queue_actions(Actions, State2#ch{queue_states = QueueStates}),
             case rabbit_event:stats_level(State, #ch.stats_timer) of
                 fine ->
                     ?INCR_STATS(exchange_stats, XName, 1, publish),
@@ -2165,6 +2169,17 @@ process_routing_confirm(MsgSeqNo, QRefs, XName, State)
     State#ch{unconfirmed =
              rabbit_confirms:insert(MsgSeqNo, QRefs, XName, State#ch.unconfirmed)}.
 
+notify_published_queue_types(QueueTypes,
+                             #ch{cfg = #conf{conn_pid = ConnPid},
+                                 queue_types_published = QTs0} = State) ->
+    QTs = sets:union(QTs0, sets:from_list(QueueTypes, [{version, 2}])),
+    sets:fold(
+      fun(QT, ok) ->
+              gen_server:cast(ConnPid,
+                              {channel_published_to_queue_type, self(), QT})
+      end, ok, sets:subtract(QTs, QTs0)),
+    State#ch{queue_types_published = QTs}.
+
 confirm(MsgSeqNos, QRef, State = #ch{unconfirmed = UC}) ->
     %% NOTE: if queue name does not exist here it's likely that the ref also
     %% does not exist in unconfirmed messages.

deps/rabbit/src/rabbit_reader.erl

@@ -113,7 +113,14 @@
   %% a set of the reasons why we are
   %% blocked: {resource, memory}, {resource, disk}.
   %% More reasons can be added in the future.
-  blocked_by,
+  blocked_by = sets:new([{version, 2}]) ::
+               sets:set(flow | {resource,
+                                rabbit_alarm:resource_alarm_source()}),
+  %% the set of queue types which have been published to
+  %% by channels on this connection, used for per-queue
+  %% type disk alarm blocking
+  queue_types_published = #{} :: #{ChannelPid :: pid() =>
+                                   sets:set(rabbit_queue_type:queue_type())},
   %% true if received any publishes, false otherwise
   %% note that this will also be true when connection is
   %% already blocked
@@ -335,7 +342,6 @@ start_connection(Parent, HelperSups, RanchRef, Deb, Sock) ->
                 throttle            = #throttle{
                                          last_blocked_at = never,
                                          should_block = false,
-                                         blocked_by = sets:new([{version, 2}]),
                                          connection_blocked_message_sent = false
                                          },
                 proxy_socket = rabbit_net:maybe_get_proxy_socket(Sock)},
@@ -677,6 +683,14 @@ handle_other({'$gen_cast', {force_event_refresh, Ref}}, State)
 handle_other({'$gen_cast', {force_event_refresh, _Ref}}, State) ->
     %% Ignore, we will emit a created event once we start running.
     State;
+handle_other({'$gen_cast', {channel_published_to_queue_type, ChPid, QT}},
+             #v1{throttle = Throttle0} = State0) ->
+    QTs = maps:update_with(
+            ChPid, fun(ChQTs) -> sets:add_element(QT, ChQTs) end,
+            sets:from_list([QT], [{version, 2}]),
+            Throttle0#throttle.queue_types_published),
+    Throttle = Throttle0#throttle{queue_types_published = QTs},
+    control_throttle(State0#v1{throttle = Throttle});
 handle_other(ensure_stats, State) ->
     ensure_stats_timer(State);
 handle_other(emit_stats, State) ->
@@ -1007,14 +1021,21 @@ is_over_node_channel_limit() ->
             end
     end.
 
-channel_cleanup(ChPid, State = #v1{channel_count = ChannelCount}) ->
+channel_cleanup(ChPid, State = #v1{channel_count = ChannelCount,
+                                   throttle = Throttle0} = State) ->
     case get({ch_pid, ChPid}) of
-        undefined       -> {undefined, State};
-        {Channel, MRef} -> credit_flow:peer_down(ChPid),
-                           erase({channel, Channel}),
-                           erase({ch_pid, ChPid}),
-                           erlang:demonitor(MRef, [flush]),
-                           {Channel, State#v1{channel_count = ChannelCount - 1}}
+        undefined ->
+            {undefined, State};
+        {Channel, MRef} ->
+            credit_flow:peer_down(ChPid),
+            erase({channel, Channel}),
+            erase({ch_pid, ChPid}),
+            erlang:demonitor(MRef, [flush]),
+            QT = maps:remove(ChPid,
+                             Throttle0#throttle.queue_types_published),
+            Throttle = Throttle0#throttle{queue_types_published = QT},
+            {Channel, State#v1{channel_count = ChannelCount - 1,
+                               throttle = Throttle}}
     end.
 
 all_channels() -> [ChPid || {{ch_pid, ChPid}, _ChannelMRef} <- get()].
@@ -1738,22 +1759,44 @@ update_last_blocked_at(Throttle) ->
 connection_blocked_message_sent(
     #throttle{connection_blocked_message_sent = BS}) -> BS.
 
-should_send_blocked(Throttle = #throttle{blocked_by = Reasons}) ->
+should_send_blocked(Throttle) ->
     should_block(Throttle)
     andalso
-    sets:size(sets:del_element(flow, Reasons)) =/= 0
+    do_throttle_reasons_apply(Throttle)
     andalso
     not connection_blocked_message_sent(Throttle).
 
-should_send_unblocked(Throttle = #throttle{blocked_by = Reasons}) ->
+should_send_unblocked(Throttle) ->
     connection_blocked_message_sent(Throttle)
     andalso
-    sets:size(sets:del_element(flow, Reasons)) == 0.
+    not do_throttle_reasons_apply(Throttle).
+
+do_throttle_reasons_apply(#throttle{blocked_by = Reasons} = Throttle) ->
+    lists:any(
+      fun ({resource, disk}) ->
+              true;
+          ({resource, memory}) ->
+              true;
+          ({resource, {disk, QT}}) ->
+              has_published_to_queue_type(QT, Throttle);
+          (_) ->
+              %% Flow control should not send connection.blocked
+              false
+      end, sets:to_list(Reasons)).
+
+has_published_to_queue_type(QT, #throttle{queue_types_published = QTs}) ->
+    rabbit_misc:maps_any(
+      fun(_ChPid, ChQT) -> sets:is_element(QT, ChQT) end, QTs).
 
 %% Returns true if we have a reason to block
 %% this connection.
-has_reasons_to_block(#throttle{blocked_by = Reasons}) ->
-    sets:size(Reasons) > 0.
+has_reasons_to_block(#throttle{blocked_by = Reasons} = Throttle) ->
+    lists:any(
+      fun ({resource, {disk, QType}}) ->
+              has_published_to_queue_type(QType, Throttle);
+          (_) ->
+              true
+      end, sets:to_list(Reasons)).
 
 is_blocked_by_flow(#throttle{blocked_by = Reasons}) ->
     sets:is_element(flow, Reasons).