AMQP 1.0: Handle per-queue-type disk alarms

the-mikedavis · the-mikedavis · commit d8b19d358259 · 2025-11-11T14:25:22.000-05:00
diff --git a/deps/rabbit/src/rabbit_amqp_session.erl b/deps/rabbit/src/rabbit_amqp_session.erl
@@ -381,6 +381,8 @@
           stashed_eol = [] :: [rabbit_amqqueue:name()],
 
           queue_states = rabbit_queue_type:init() :: rabbit_queue_type:state(),
+          queue_types_published = sets:new([{version, 2}]) ::
+                                  sets:set(rabbit_queue_type:queue_type()),
           permission_cache = [] :: permission_cache(),
           topic_permission_cache = [] :: topic_permission_cache()
          }).
@@ -448,9 +450,14 @@ init({ReaderPid, WriterPid, ChannelNum, MaxFrameSize, User, Vhost, ContainerId,
     true = is_valid_max(MaxLinkCredit),
     true = is_valid_max(MaxQueueCredit),
     true = is_valid_max(MaxIncomingWindow),
-    IncomingWindow = case sets:is_empty(Alarms) of
-                         true -> MaxIncomingWindow;
-                         false -> 0
+    InResourceAlarm = sets:fold(fun ({disk, _}, Acc) ->
+                                        Acc;
+                                    (_, _) ->
+                                        true
+                                end, false, Alarms),
+    IncomingWindow = case InResourceAlarm of
+                         true -> 0;
+                         false -> MaxIncomingWindow
                      end,
     NextOutgoingId = ?INITIAL_OUTGOING_TRANSFER_ID,
 
@@ -582,36 +589,17 @@ handle_cast({queue_event, _, _} = QEvent, State0) ->
               log_error_and_close_session(Error, State0)
     end;
 handle_cast({conserve_resources, Alarm, Conserve},
-            #state{incoming_window = IncomingWindow0,
-                   cfg = #cfg{resource_alarms = Alarms0,
-                              incoming_window_margin = Margin0,
+            #state{cfg = #cfg{resource_alarms = Alarms0,
                               writer_pid = WriterPid,
-                              channel_num = Ch,
-                              max_incoming_window = MaxIncomingWindow
+                              channel_num = Ch
                              } = Cfg
                   } = State0) ->
     Alarms = case Conserve of
                  true -> sets:add_element(Alarm, Alarms0);
                  false -> sets:del_element(Alarm, Alarms0)
              end,
-    {SendFlow, IncomingWindow, Margin} =
-    case {sets:is_empty(Alarms0), sets:is_empty(Alarms)} of
-        {true, false} ->
-            %% Alarm kicked in.
-            %% Notify the client to not send us any more TRANSFERs. Since we decrase
-            %% our incoming window dynamically, there might be incoming in-flight
-            %% TRANSFERs. So, let's be lax and allow for some excess TRANSFERs.
-            {true, 0, MaxIncomingWindow};
-        {false, true} ->
-            %% All alarms cleared.
-            %% Notify the client that it can resume sending us TRANSFERs.
-            {true, MaxIncomingWindow, 0};
-        _ ->
-            {false, IncomingWindow0, Margin0}
-    end,
-    State = State0#state{incoming_window = IncomingWindow,
-                         cfg = Cfg#cfg{resource_alarms = Alarms,
-                                       incoming_window_margin = Margin}},
+    State1 = State0#state{cfg = Cfg#cfg{resource_alarms = Alarms}},
+    {SendFlow, State} = check_resource_alarm(State0, State1),
     case SendFlow of
         true ->
             Flow = session_flow_fields(#'v1_0.flow'{}, State),
@@ -637,6 +625,41 @@ handle_cast({reset_authz, User}, #state{cfg = Cfg} = State0) ->
 handle_cast(shutdown, State) ->
     {stop, normal, State}.
 
+is_in_resource_alarm(#state{cfg = #cfg{resource_alarms = Alarms},
+                            queue_types_published = QTs}) ->
+    sets:fold(
+      fun ({disk, QT}, Acc) ->
+              Acc orelse sets:is_element(QT, QTs);
+          (_, _) ->
+              true
+      end, false, Alarms).
+
+check_resource_alarm(State0,
+                     #state{incoming_window = IncomingWindow0,
+                            cfg = #cfg{incoming_window_margin = Margin0,
+                                       max_incoming_window = MaxIncomingWindow
+                                      } = Cfg} = State1) ->
+    WasBlocked = is_in_resource_alarm(State0),
+    IsBlocked = is_in_resource_alarm(State1),
+    {SendFlow, IncomingWindow, Margin} =
+    case IsBlocked of
+        true when not WasBlocked ->
+            %% Alarm kicked in.
+            %% Notify the client to not send us any more TRANSFERs. Since we decrase
+            %% our incoming window dynamically, there might be incoming in-flight
+            %% TRANSFERs. So, let's be lax and allow for some excess TRANSFERs.
+            {true, 0, MaxIncomingWindow};
+        false when WasBlocked ->
+            %% All alarms cleared.
+            %% Notify the client that it can resume sending us TRANSFERs.
+            {true, MaxIncomingWindow, 0};
+        _ ->
+            {false, IncomingWindow0, Margin0}
+    end,
+    State = State1#state{incoming_window = IncomingWindow,
+                         cfg = Cfg#cfg{incoming_window_margin = Margin}},
+    {SendFlow, State}.
+
 log_error_and_close_session(
   Error, State = #state{cfg = #cfg{reader_pid = ReaderPid,
                                    writer_pid = WriterPid,
@@ -1940,7 +1963,6 @@ session_flow_control_received_transfer(
          incoming_window = InWindow0,
          remote_outgoing_window = RemoteOutgoingWindow,
          cfg = #cfg{incoming_window_margin = Margin,
-                    resource_alarms = Alarms,
                     max_incoming_window = MaxIncomingWindow}
         } = State) ->
     InWindow1 = InWindow0 - 1,
@@ -1954,7 +1976,7 @@ session_flow_control_received_transfer(
             ok
     end,
     {Flows, InWindow} = case InWindow1 =< (MaxIncomingWindow div 2) andalso
-                             sets:is_empty(Alarms) of
+                             not is_in_resource_alarm(State) of
                             true ->
                                 %% We've reached halfway and there are no
                                 %% disk or memory alarm, open the window.
@@ -2371,6 +2393,7 @@ incoming_link_transfer(
                  multi_transfer_msg = MultiTransfer
                 } = Link0,
   State0 = #state{queue_states = QStates0,
+                  queue_types_published = QTs0,
                   permission_cache = PermCache0,
                   topic_permission_cache = TopicPermCache0,
                   cfg = #cfg{user = User = #user{username = Username},
@@ -2414,19 +2437,26 @@ incoming_link_transfer(
             Qs = rabbit_amqqueue:prepend_extra_bcc(Qs0),
             case rabbit_queue_type:deliver(Qs, Mc, Opts, QStates0) of
                 {ok, QStates, Actions} ->
+                    QTs1 = sets:from_list(rabbit_amqqueue:queue_types(Qs),
+                                          [{version, 2}]),
+                    QTs = sets:union(QTs0, QTs1),
                     State1 = State0#state{queue_states = QStates,
+                                          queue_types_published = QTs,
                                           permission_cache = PermCache,
                                           topic_permission_cache = TopicPermCache},
                     %% Confirms must be registered before processing actions
                     %% because actions may contain rejections of publishes.
                     {U, Reply0} = process_routing_confirm(
                                     Qs, Settled, DeliveryId, U0),
-                    State = handle_queue_actions(Actions, State1),
+                    State2 = handle_queue_actions(Actions, State1),
+                    {SendAlarmFlow, State} = check_resource_alarm(
+                                               State0, State2),
                     DeliveryCount = add(DeliveryCount0, 1),
                     Credit1 = Credit0 - 1,
                     {Credit, Reply1} = maybe_grant_link_credit(
                                          Credit1, MaxLinkCredit,
-                                         DeliveryCount, map_size(U), Handle),
+                                         DeliveryCount, map_size(U), Handle,
+                                         SendAlarmFlow),
                     Reply = Reply0 ++ Reply1,
                     Link = Link0#incoming_link{
                              delivery_count = DeliveryCount,
@@ -2461,7 +2491,8 @@ incoming_link_transfer(
                     Credit1 = Credit0 - 1,
                     {Credit, Reply0} = maybe_grant_link_credit(
                                          Credit1, MaxLinkCredit,
-                                         DeliveryCount, map_size(U0), Handle),
+                                         DeliveryCount, map_size(U0), Handle,
+                                         false),
                     Reply = [Disposition | Reply0],
                     Link = Link0#incoming_link{
                              delivery_count = DeliveryCount,
@@ -2574,8 +2605,9 @@ rejected(DeliveryId, Error) ->
                         settled = true,
                         state = #'v1_0.rejected'{error = Error}}.
 
-maybe_grant_link_credit(Credit, MaxLinkCredit, DeliveryCount, NumUnconfirmed, Handle) ->
-    case grant_link_credit(Credit, MaxLinkCredit, NumUnconfirmed) of
+maybe_grant_link_credit(Credit, MaxLinkCredit, DeliveryCount, NumUnconfirmed,
+                        Handle, AlarmFlow) ->
+    case grant_link_credit(Credit, MaxLinkCredit, NumUnconfirmed) orelse AlarmFlow of
         true ->
             {MaxLinkCredit, [flow(Handle, DeliveryCount, MaxLinkCredit)]};
         false ->
diff --git a/deps/rabbit/test/amqp_client_SUITE.erl b/deps/rabbit/test/amqp_client_SUITE.erl
@@ -116,6 +116,7 @@ groups() ->
        resource_alarm_before_session_begin,
        resource_alarm_after_session_begin,
        resource_alarm_send_many,
+       per_queue_type_disk_alarm,
        max_message_size_client_to_server,
        max_message_size_server_to_client,
        global_counters,
@@ -225,7 +226,13 @@ init_per_suite(Config) ->
     rabbit_ct_helpers:log_environment(),
     rabbit_ct_helpers:merge_app_env(
       Config, {rabbit, [{quorum_tick_interval, 1000},
-                        {stream_tick_interval, 1000}
+                        {stream_tick_interval, 1000},
+                        %% Imaginary mount-point for per-queue-type disk alarms
+                        {disk_free_limits,
+                         #{1 => #{name => <<"streaming">>,
+                                  mount => "/does/not/exist",
+                                  limit => "2GB",
+                                  queue_types => [<<"stream">>]}}}
                        ]}).
 
 end_per_suite(Config) ->
@@ -3263,6 +3270,76 @@ auth_attempt_metrics(Config) ->
     ?assertEqual(0, proplists:get_value(auth_attempts_failed, Attempt2)),
     ?assertEqual(1, proplists:get_value(auth_attempts_succeeded, Attempt2)).
 
+per_queue_type_disk_alarm(Config) ->
+    Prefix = atom_to_binary(?FUNCTION_NAME),
+    Resource = {disk, rabbit_stream_queue},
+    CQ = <<Prefix/binary, "-classic">>,
+    SQ = <<Prefix/binary, "-stream">>,
+    {Conn, Ch} = rabbit_ct_client_helpers:open_connection_and_channel(Config),
+    #'queue.declare_ok'{} = amqp_channel:call(Ch, #'queue.declare'{queue = CQ}),
+    #'queue.declare_ok'{} = amqp_channel:call(
+                              Ch, #'queue.declare'{
+                                     queue = SQ,
+                                     durable = true,
+                                     arguments = [{<<"x-queue-type">>, longstr, <<"stream">>}]}),
+
+    OpnConf = connection_config(Config),
+    {ok, Connection} = amqp10_client:open_connection(OpnConf),
+
+    %% Set the alarm for the stream queue type.
+    ok = rabbit_ct_broker_helpers:set_alarm(Config, 0, Resource),
+
+    %% Attach one sender to the CQ and one to the SQ.
+    {ok, Session1} = amqp10_client:begin_session_sync(Connection),
+    {ok, Sender1} = amqp10_client:attach_sender_link(
+                      Session1, <<Prefix/binary, "-cq-sender">>,
+                      rabbitmq_amqp_address:queue(CQ), unsettled),
+    {ok, Session2} = amqp10_client:begin_session_sync(Connection),
+    {ok, Sender2} = amqp10_client:attach_sender_link(
+                      Session2, <<Prefix/binary, "-sq-sender">>,
+                      rabbitmq_amqp_address:queue(SQ), unsettled),
+
+    %% Both senders initially have link and session credit.
+    ok = wait_for_credit(Sender1),
+    ok = wait_for_credit(Sender2),
+    Tag1 = <<"tag1">>,
+    Msg1 = amqp10_msg:new(Tag1, <<"m1">>, false),
+    ?assertEqual(ok,
+                 amqp10_client:send_msg(Sender1, Msg1)),
+    ok = wait_for_accepted(Tag1),
+    ?assertEqual(ok,
+                 amqp10_client:send_msg(Sender2, Msg1)),
+    ok = wait_for_accepted(Tag1),
+
+    %% Once the SQ sender has delivered to a stream, it becomes blocked by
+    %% session flow control.
+    Tag2 = <<"tag2">>,
+    Msg2 = amqp10_msg:new(Tag2, <<"m2">>, false),
+    ?assertEqual(ok,
+                 amqp10_client:send_msg(Sender1, Msg2)),
+    ok = wait_for_accepted(Tag2),
+    ?assertEqual({error, remote_incoming_window_exceeded},
+                 amqp10_client:send_msg(Sender2, Msg2)),
+
+    %% Clear the alarm and the SQ sender can then send transfers.
+    ok = rabbit_ct_broker_helpers:clear_alarm(Config, 0, Resource),
+    Tag3 = <<"tag3">>,
+    Msg3 = amqp10_msg:new(Tag3, <<"m3">>, false),
+    ?assertEqual(ok,
+                 amqp10_client:send_msg(Sender1, Msg3)),
+    ok = wait_for_accepted(Tag3),
+    ?assertEqual(ok,
+                 amqp10_client:send_msg(Sender2, Msg3)),
+    ok = wait_for_accepted(Tag3),
+
+    ok = amqp10_client:detach_link(Sender1),
+    ok = end_session_sync(Session1),
+    ok = amqp10_client:detach_link(Sender2),
+    ok = end_session_sync(Session2),
+    #'queue.delete_ok'{} = amqp_channel:call(Ch, #'queue.delete'{queue = CQ}),
+    #'queue.delete_ok'{} = amqp_channel:call(Ch, #'queue.delete'{queue = SQ}),
+    ok = rabbit_ct_client_helpers:close_connection_and_channel(Conn, Ch).
+
 max_message_size_client_to_server(Config) ->
     DefaultMaxMessageSize = rpc(Config, persistent_term, get, [max_message_size]),
     %% Limit the server to only accept messages up to 2KB.
diff --git a/deps/rabbitmq_ct_helpers/src/rabbit_ct_broker_helpers.erl b/deps/rabbitmq_ct_helpers/src/rabbit_ct_broker_helpers.erl
@@ -1815,6 +1815,8 @@ set_alarm(Config, Node, file_descriptor_limit = Resource) ->
 set_alarm(Config, Node, memory = Resource) ->
     rpc(Config, Node, rabbit_alarm, set_alarm, [{{resource_limit, Resource, Node}, []}]);
 set_alarm(Config, Node, disk = Resource) ->
+    rpc(Config, Node, rabbit_alarm, set_alarm, [{{resource_limit, Resource, Node}, []}]);
+set_alarm(Config, Node, {disk, QueueType} = Resource) when is_atom(QueueType) ->
     rpc(Config, Node, rabbit_alarm, set_alarm, [{{resource_limit, Resource, Node}, []}]).
 
 get_alarms(Config, Node) ->
@@ -1828,6 +1830,8 @@ clear_alarm(Config, Node, file_descriptor_limit = Resource) ->
 clear_alarm(Config, Node, memory = Resource) ->
     rpc(Config, Node, rabbit_alarm, clear_alarm, [{resource_limit, Resource, Node}]);
 clear_alarm(Config, Node, disk = Resource) ->
+    rpc(Config, Node, rabbit_alarm, clear_alarm, [{resource_limit, Resource, Node}]);
+clear_alarm(Config, Node, {disk, QueueType} = Resource) when is_atom(QueueType) ->
     rpc(Config, Node, rabbit_alarm, clear_alarm, [{resource_limit, Resource, Node}]).
 
 clear_all_alarms(Config, Node) ->
