amqp_client: Per-channel queue_types_published tracking for direct connections

Previously, queue_types_published in amqp_gen_connection was a flat set
with no per-channel cleanup. Once a connection published to a queue type
it was permanently considered as publishing to that type, causing
over-blocking when channels closed.

Fix by threading the server channel pid through the channel startup path
as a tagged tuple {direct, ServerPid} | {network, WriterPid}:

- amqp_channel_sup: start_writer/6 returns a tagged writer tuple;
  start_link/6 passes it through to the caller
- amqp_channel: set_writer/2 unwraps the tagged tuple
- amqp_channels_manager: stores ClientPid -> TaggedWriter mapping;
  sends channel_closed/2 with the tagged writer on channel exit
- amqp_gen_connection: tracks queue_types_published per server channel
  pid (keyed by {direct, ServerPid}); removes the entry on channel_closed

Network connections are unaffected: rabbit_reader handles blocking
independently and amqp_gen_connection ignores {network, _} channel_closed.

Test: amqp_client_SUITE:per_queue_type_disk_alarm_direct_connection
verifies that a direct connection is unblocked when the publishing
channel closes, even while the alarm remains active.

Author: lukebakken

deps/amqp_client/src/amqp_channel.erl

@@ -352,7 +352,11 @@ start_link(Driver, Connection, ChannelNumber, Consumer, Identity) ->
     gen_server:start_link(
       ?MODULE, [Driver, Connection, ChannelNumber, Consumer, Identity], []).
 
-set_writer(Pid, Writer) ->
+set_writer(Pid, {network, Writer}) ->
+    set_writer(Pid, Writer);
+set_writer(Pid, {direct, Writer}) ->
+    set_writer(Pid, Writer);
+set_writer(Pid, Writer) when is_pid(Pid) andalso is_pid(Writer) ->
     gen_server:cast(Pid, {set_writer, Writer}).
 
 enable_delivery_flow_control(Pid) ->

deps/amqp_client/src/amqp_channel_sup.erl

@@ -34,10 +34,10 @@ start_link(Type, Connection, ConnName, InfraArgs, ChNumber,
                   modules => [amqp_channel]},
     {ok, ChPid} = supervisor:start_child(Sup, ChildSpec),
     case start_writer(Sup, Type, InfraArgs, ConnName, ChNumber, ChPid) of
-        {ok, Writer} ->
-            amqp_channel:set_writer(ChPid, Writer),
+        {ok, TaggedWriter} ->
+            amqp_channel:set_writer(ChPid, TaggedWriter),
             {ok, AState} = init_command_assembler(Type),
-            {ok, Sup, {ChPid, AState}};
+            {ok, Sup, {ChPid, TaggedWriter, AState}};
         {error, _}=Error ->
             rabbit_misc:shutdown_supervisor(Sup),
             Error
@@ -59,12 +59,11 @@ start_writer(_Sup, direct, [ConnPid, Node, User, VHost, Collector, AmqpParams],
     case rpc:call(Node, rabbit_direct, start_channel,
                [ChNumber, ChPid, ConnPid, ConnName, ?PROTOCOL, User,
                 VHost, ?CLIENT_CAPABILITIES, Collector, AmqpParams], ?DIRECT_OPERATION_TIMEOUT) of
-        {ok, _Writer} = Reply ->
-            Reply;
+        {ok, Writer} ->
+            {ok, {direct, Writer}};
         {badrpc, Reason} ->
             {error, {Reason, Node}};
-        Error ->
-            Error
+        Error -> Error
     end;
 start_writer(Sup, network, [Sock, FrameMax], ConnName, ChNumber, ChPid) ->
     GCThreshold = application:get_env(amqp_client, writer_gc_threshold, ?DEFAULT_GC_THRESHOLD),
@@ -77,7 +76,11 @@ start_writer(Sup, network, [Sock, FrameMax], ConnName, ChNumber, ChPid) ->
                   shutdown => ?WORKER_WAIT,
                   type => worker,
                   modules => [rabbit_writer]},
-    supervisor:start_child(Sup, ChildSpec).
+    case supervisor:start_child(Sup, ChildSpec) of
+        {ok, Writer} ->
+            {ok, {network, Writer}};
+        Error -> Error
+    end.
 
 init_command_assembler(direct)  -> {ok, none};
 init_command_assembler(network) -> rabbit_command_assembler:init().

deps/amqp_client/src/amqp_channels_manager.erl

@@ -23,6 +23,7 @@
                 channel_sup_sup,
                 map_num_pa      = gb_trees:empty(), %% Number -> {Pid, AState}
                 map_pid_num     = #{},       %% Pid -> Number
+                map_pid_writer  = #{},       %% Pid -> {direct|network, WriterPid}
                 channel_max     = ?MAX_CHANNEL_NUMBER,
                 closing         = false}).
 
@@ -107,8 +108,8 @@ handle_open_channel(ProposedNumber, Consumer, InfraArgs,
         {ok, Number} ->
             case amqp_channel_sup_sup:start_channel_sup(ChSupSup, InfraArgs,
                                                         Number, Consumer) of
-                {ok, _ChSup, {Ch, AState}} ->
-                    NewState = internal_register(Number, Ch, AState, State),
+                {ok, _ChSup, {Ch, TaggedWriter, AState}} ->
+                    NewState = internal_register(Number, Ch, TaggedWriter, AState, State),
                     erlang:monitor(process, Ch),
                     {reply, {ok, Ch}, NewState};
                 {error, _} = Error ->
@@ -158,12 +159,15 @@ handle_down(Pid, Reason, State) ->
         Number    -> handle_channel_down(Pid, Number, Reason, State)
     end.
 
-handle_channel_down(Pid, Number, Reason, State) ->
+handle_channel_down(Pid, Number, Reason, State = #state{connection = Connection,
+                                                        map_pid_writer = MapPW}) ->
     maybe_report_down(Pid, case Reason of {shutdown, R} -> R;
                                           _             -> Reason
                            end,
                       State),
+    TaggedWriter = maps:get(Pid, MapPW),
     NewState = internal_unregister(Number, Pid, State),
+    amqp_gen_connection:channel_closed(Connection, TaggedWriter),
     check_all_channels_terminated(NewState),
     {noreply, NewState}.
 
@@ -214,19 +218,25 @@ internal_pass_frame(Number, Frame, State) ->
             internal_update_npa(Number, ChPid, NewAState, State)
     end.
 
-internal_register(Number, Pid, AState,
-                  State = #state{map_num_pa = MapNPA, map_pid_num = MapPN}) ->
+internal_register(Number, Pid, TaggedWriter, AState,
+                  State = #state{map_num_pa = MapNPA, map_pid_num = MapPN,
+                                 map_pid_writer = MapPW}) ->
     MapNPA1 = gb_trees:enter(Number, {Pid, AState}, MapNPA),
     MapPN1 = maps:put(Pid, Number, MapPN),
-    State#state{map_num_pa  = MapNPA1,
-                map_pid_num = MapPN1}.
+    MapPW1 = maps:put(Pid, TaggedWriter, MapPW),
+    State#state{map_num_pa     = MapNPA1,
+                map_pid_num    = MapPN1,
+                map_pid_writer = MapPW1}.
 
 internal_unregister(Number, Pid,
-                    State = #state{map_num_pa = MapNPA, map_pid_num = MapPN}) ->
+                    State = #state{map_num_pa = MapNPA, map_pid_num = MapPN,
+                                   map_pid_writer = MapPW}) ->
     MapNPA1 = gb_trees:delete(Number, MapNPA),
     MapPN1 = maps:remove(Pid, MapPN),
-    State#state{map_num_pa  = MapNPA1,
-                map_pid_num = MapPN1}.
+    MapPW1 = maps:remove(Pid, MapPW),
+    State#state{map_num_pa     = MapNPA1,
+                map_pid_num    = MapPN1,
+                map_pid_writer = MapPW1}.
 
 internal_is_empty(#state{map_num_pa = MapNPA}) ->
     gb_trees:is_empty(MapNPA).

deps/amqp_client/src/amqp_gen_connection.erl

@@ -14,9 +14,9 @@
 -behaviour(gen_server).
 
 -export([start_link/2, connect/1, open_channel/3, hard_error_in_channel/3,
-         channel_internal_error/3, server_misbehaved/2, channels_terminated/1,
-         close/3, server_close/2, info/2, info_keys/0, info_keys/1,
-         register_blocked_handler/2, update_secret/2]).
+         channel_internal_error/3, channel_closed/2, server_misbehaved/2,
+         channels_terminated/1, close/3, server_close/2, info/2, info_keys/0,
+         info_keys/1, register_blocked_handler/2, update_secret/2]).
 -export([init/1, terminate/2, code_change/3, handle_call/3, handle_cast/2,
          handle_info/2]).
 
@@ -32,7 +32,7 @@
                 %% connection.block, connection.unblock handler
                 block_handler,
                 blocked_by = sets:new([{version, 2}]),
-                queue_types_published = sets:new([{version, 2}]),
+                queue_types_published = #{},
                 closing = false %% #closing{} | false
                }).
 
@@ -97,6 +97,9 @@ hard_error_in_channel(Pid, ChannelPid, Reason) ->
 channel_internal_error(Pid, ChannelPid, Reason) ->
     gen_server:cast(Pid, {channel_internal_error, ChannelPid, Reason}).
 
+channel_closed(Pid, TaggedWriter) ->
+    gen_server:cast(Pid, {channel_closed, TaggedWriter}).
+
 server_misbehaved(Pid, AmqpError) ->
     gen_server:cast(Pid, {server_misbehaved, AmqpError}).
 
@@ -223,10 +226,23 @@ handle_cast({conserve_resources, Source, Conserve},
                  end,
     State1 = State#state{blocked_by = BlockedBy1},
     maybe_block(State, State1);
-handle_cast({channel_published_to_queue_type, _ChPid, QT},
+handle_cast({channel_published_to_queue_type, ServerChPid, QT},
+            #state{queue_types_published = QTs} = State) ->
+    %% ServerChPid is the server-side channel pid. We key by {direct, ServerChPid}
+    %% so channel_closed can remove the entry using the tagged writer.
+    Key = {direct, ServerChPid},
+    ChQTs = maps:get(Key, QTs, sets:new([{version, 2}])),
+    State1 = State#state{queue_types_published =
+                             QTs#{Key => sets:add_element(QT, ChQTs)}},
+    maybe_block(State, State1);
+handle_cast({channel_closed, {direct, ServerChPid}},
             #state{queue_types_published = QTs} = State) ->
-    State1 = State#state{queue_types_published = sets:add_element(QT, QTs)},
-    maybe_block(State, State1).
+    State1 = State#state{queue_types_published =
+                             maps:remove({direct, ServerChPid}, QTs)},
+    maybe_block(State, State1);
+handle_cast({channel_closed, {network, _WriterPid}}, State) ->
+    %% Network connections are handled by rabbit_reader, not here.
+    {noreply, State}.
 
 %% @private
 handle_info({'DOWN', _, process, BlockHandler, Reason},
@@ -284,7 +300,9 @@ maybe_block(State0, State1) ->
 
 should_block(#state{blocked_by = BlockedBy, queue_types_published = QTs}) ->
     lists:any(fun ({disk, QT}) ->
-                      sets:is_element(QT, QTs);
+                      maps:fold(fun(_, ChQTs, Acc) ->
+                                        Acc orelse sets:is_element(QT, ChQTs)
+                                end, false, QTs);
                   (_Resource) ->
                       true
               end, sets:to_list(BlockedBy)).

deps/rabbit/test/amqp_client_SUITE.erl

@@ -3824,7 +3824,8 @@ per_queue_type_disk_alarm(Config) ->
 per_queue_type_disk_alarm_direct_connection(Config) ->
     %% Test that a direct connection (used by plugins such as STOMP) is
     %% blocked when a per-queue-type disk alarm fires and the connection
-    %% has published to that queue type, and unblocked when the alarm clears.
+    %% has published to that queue type, and unblocked when the publishing
+    %% channel closes - even while the alarm remains active.
     Node = rabbit_ct_broker_helpers:get_node_config(Config, 0, nodename),
     Resource = {disk, rabbit_classic_queue},
     QName = atom_to_binary(?FUNCTION_NAME),
@@ -3851,15 +3852,18 @@ per_queue_type_disk_alarm_direct_connection(Config) ->
     after 5000 -> exit(connection_was_not_blocked)
     end,
 
-    %% Clear the alarm. The connection should unblock.
-    ok = rpc:call(Node, rabbit_alarm, clear_alarm, [{resource_limit, Resource, Node}]),
+    %% Close the publishing channel. The connection should unblock since no
+    %% active channel has published to the alarmed queue type.
+    ok = amqp_channel:close(Ch),
     receive
         #'connection.unblocked'{} -> ok
     after 5000 -> exit(connection_was_not_unblocked)
     end,
 
-    #'queue.delete_ok'{} = amqp_channel:call(Ch, #'queue.delete'{queue = QName}),
-    ok = amqp_channel:close(Ch),
+    ok = rpc:call(Node, rabbit_alarm, clear_alarm, [{resource_limit, Resource, Node}]),
+    {ok, Ch2} = amqp_connection:open_channel(Conn),
+    #'queue.delete_ok'{} = amqp_channel:call(Ch2, #'queue.delete'{queue = QName}),
+    ok = amqp_channel:close(Ch2),
     ok = amqp_connection:close(Conn).
 
 max_message_size_client_to_server(Config) ->