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,18 +215,30 @@ handle_cast({register_blocked_handler, HandlerPid},
     {noreply, State1};
 handle_cast({conserve_resources, Source, Conserve},
             #state{blocked_by = BlockedBy} = State) ->
-    WasNotBlocked = sets:is_empty(BlockedBy),
+    WasBlocked = should_block(State),
     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
+    case should_block(State1) of
         true ->
             handle_method(#'connection.unblocked'{}, State1);
-        false when WasNotBlocked ->
+        false when not WasBlocked ->
+            handle_method(#'connection.blocked'{}, State1);
+        false ->
+            {noreply, State1}
+    end;
+handle_cast({channel_published_to_queue_type, _ChPid, QT},
+            #state{queue_types_published = QTs} = State) ->
+    WasBlocked = should_block(State),
+    State1 = State#state{queue_types_published = sets:add_element(QT, QTs)},
+    case should_block(State1) of
+        true ->
+            handle_method(#'connection.unblocked'{}, State1);
+        false when not WasBlocked ->
             handle_method(#'connection.blocked'{}, State1);
         false ->
             {noreply, State1}
@@ -274,6 +287,13 @@ 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}).
 
+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_channel.erl

@@ -170,7 +170,9 @@
              interceptor_state,
              queue_states,
              tick_timer,
-             publishing_mode = false :: boolean()
+             publishing_mode = false :: boolean(),
+             queue_types_published = sets:new([{version, 2}]) ::
+                                     sets:set(rabbit_queue_type:queue_type())
             }).
 
 -define(QUEUE, lqueue).
@@ -2097,9 +2099,10 @@ deliver_to_queues(XName,
             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(Qs, 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 +2168,27 @@ process_routing_confirm(MsgSeqNo, QRefs, XName, State)
     State#ch{unconfirmed =
              rabbit_confirms:insert(MsgSeqNo, QRefs, XName, State#ch.unconfirmed)}.
 
+notify_published_queue_types(Qs, #ch{cfg = #conf{conn_pid = ConnPid},
+                                     queue_types_published = QTs0} = State0) ->
+    QTs = lists:foldl(
+            fun(Q0, Acc) ->
+                    Q = case Q0 of
+                            {Q1, _RouteInfo} -> Q1;
+                            _ -> Q0
+                        end,
+                    QT = amqqueue:get_type(Q),
+                    case sets:is_element(QT, Acc) of
+                        true ->
+                            Acc;
+                        false ->
+                            gen_server:cast(
+                              ConnPid,
+                              {channel_published_to_queue_type, self(), QT}),
+                            sets:add_element(QT, Acc)
+                    end
+            end, QTs0, Qs),
+    State0#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).