Add channel limit per node

SimonUnge · michaelklishin · commit 34a3796d2944 · 2024-01-31T14:30:57.000-05:00
diff --git a/deps/rabbit/priv/schema/rabbit.schema b/deps/rabbit/priv/schema/rabbit.schema
@@ -889,6 +889,9 @@ end}.
 
 {mapping, "channel_max", "rabbit.channel_max", [{datatype, integer}]}.
 
+{mapping, "channel_max_per_node", "rabbit.channel_max_per_node",
+          [{datatype, integer}, {validators, ["non_negative_integer"]}]}.
+
 %% Set the max permissible number of client connections per node.
 %% `infinity` means "no limit".
 %%
diff --git a/deps/rabbit/src/rabbit_reader.erl b/deps/rabbit/src/rabbit_reader.erl
@@ -904,7 +904,7 @@ create_channel(Channel,
                                    capabilities = Capabilities,
                                    user = #user{username = Username} = User}
                    } = State) ->
-    case rabbit_auth_backend_internal:is_over_channel_limit(Username) of
+    case is_over_limits(Username) of
         false ->
             {ok, _ChSupPid, {ChPid, AState}} =
                 rabbit_channel_sup_sup:start_channel(
@@ -915,11 +915,45 @@ create_channel(Channel,
             put({ch_pid, ChPid}, {Channel, MRef}),
             put({channel, Channel}, {ChPid, AState}),
             {ok, {ChPid, AState}, State#v1{channel_count = ChannelCount + 1}};
+        {true, Limit, Fmt} ->
+            {error, rabbit_misc:amqp_error(
+                      not_allowed,
+                      Fmt,
+                      [node(), Limit], 'none')}
+    end.
+
+is_over_limits(Username) ->
+    case rabbit_auth_backend_internal:is_over_channel_limit(Username) of
+        false ->
+            case is_over_node_channel_limit() of
+                false ->
+                    false;
+                {true, Limit} ->
+                    Fmt =
+                        "number of channels opened for node '~ts' has reached "
+                        "the maximum allowed limit of (~w)",
+                    {true, Limit, Fmt}
+            end;
         {true, Limit} ->
-            {error, rabbit_misc:amqp_error(not_allowed,
-                        "number of channels opened for user '~ts' has reached "
-                        "the maximum allowed user limit of (~w)",
-                        [Username, Limit], 'none')}
+            Fmt =
+                "number of channels opened for user '~ts' has reached "
+                "the maximum allowed user limit of (~w)",
+            {true, Limit, Fmt}
+    end.
+
+is_over_node_channel_limit() ->
+    case rabbit_misc:get_env(rabbit, channel_max_per_node, 0) of
+        0 ->
+            false;
+        NodeLimit ->
+            %% Only fetch this if a limit is set
+            CurrNodeChannels = length(rabbit_channel_tracking:list_on_node(node())),
+            case CurrNodeChannels < NodeLimit of
+                true ->
+                    false;
+                false ->
+                    {true, NodeLimit}
+            end
     end.
 
 channel_cleanup(ChPid, State = #v1{channel_count = ChannelCount}) ->
diff --git a/deps/rabbit/test/config_schema_SUITE_data/rabbit.snippets b/deps/rabbit/test/config_schema_SUITE_data/rabbit.snippets
@@ -396,6 +396,10 @@ tcp_listen_options.exit_on_close = false",
   "channel_max = 16",
   [{rabbit,[{channel_max, 16}]}],
   []},
+ {channel_max_per_node,
+  "channel_max_per_node = 16",
+  [{rabbit,[{channel_max_per_node, 16}]}],
+  []},
  {max_message_size,
   "max_message_size = 131072",
   [{rabbit, [{max_message_size, 131072}]}],
diff --git a/deps/rabbit/test/per_node_limit_SUITE.erl b/deps/rabbit/test/per_node_limit_SUITE.erl
@@ -15,15 +15,16 @@
 
 all() ->
     [
-     {group, parallel_tests}
+     {group, limit_tests}
     ].
 
 groups() ->
     [
-     {parallel_tests, [parallel], [
-                                   node_connection_limit,
-                                   vhost_limit
-                                  ]}
+     {limit_tests, [], [
+                        node_connection_limit,
+                        vhost_limit,
+                        node_channel_limit
+                       ]}
     ].
 
 suite() ->
@@ -60,9 +61,15 @@ init_per_testcase(Testcase, Config) ->
     rabbit_ct_helpers:testcase_started(Config, Testcase).
 
 end_per_testcase(vhost_limit = Testcase, Config) ->
+    set_node_limit(Config, vhost_max, infinity),
+    set_node_limit(Config, channel_max_per_node, 0),
+    set_node_limit(Config, connection_max, infinity),
     [rabbit_ct_broker_helpers:delete_vhost(Config, integer_to_binary(I)) || I <- lists:seq(1,4)],
     rabbit_ct_helpers:testcase_finished(Config, Testcase);
 end_per_testcase(Testcase, Config) ->
+    set_node_limit(Config, vhost_max, infinity),
+    set_node_limit(Config, channel_max_per_node, 0),
+    set_node_limit(Config, connection_max, infinity),
     rabbit_ct_helpers:testcase_finished(Config, Testcase).
 
 %% -------------------------------------------------------------------
@@ -71,7 +78,7 @@ end_per_testcase(Testcase, Config) ->
 
 node_connection_limit(Config) ->
     %% Set limit to 0, don't accept any connections
-    set_node_limit(Config, 0),
+    set_node_limit(Config, connection_max, 0),
     {error, not_allowed} = rabbit_ct_client_helpers:open_unmanaged_connection(Config, 0),
 
     %% Set limit to 5, accept 5 connections
@@ -80,47 +87,88 @@ node_connection_limit(Config) ->
     {error, not_allowed} = rabbit_ct_client_helpers:open_unmanaged_connection(Config, 0),
     close_all_connections(Connections),
 
-    set_node_limit(Config, infinity),
+    set_node_limit(Config, connection_max, infinity),
     C = rabbit_ct_client_helpers:open_unmanaged_connection(Config, 0),
     true = is_pid(C),
     close_all_connections([C]),
     ok.
 
 vhost_limit(Config) ->
-    set_vhost_limit(Config, 0),
+    set_node_limit(Config, vhost_max, 0),
     {'EXIT',{vhost_limit_exceeded, _}} = rabbit_ct_broker_helpers:add_vhost(Config, <<"foo">>),
 
-    set_vhost_limit(Config, 5),
+    set_node_limit(Config, vhost_max, 5),
     [ok = rabbit_ct_broker_helpers:add_vhost(Config, integer_to_binary(I)) || I <- lists:seq(1,4)],
     {'EXIT',{vhost_limit_exceeded, _}} = rabbit_ct_broker_helpers:add_vhost(Config, <<"5">>),
     [rabbit_ct_broker_helpers:delete_vhost(Config, integer_to_binary(I)) || I <- lists:seq(1,4)],
 
-    set_vhost_limit(Config, infinity),
+    set_node_limit(Config, vhost_max, infinity),
     [ok = rabbit_ct_broker_helpers:add_vhost(Config, integer_to_binary(I)) || I <- lists:seq(1,4)],
     ok = rabbit_ct_broker_helpers:add_vhost(Config, <<"5">>),
     [rabbit_ct_broker_helpers:delete_vhost(Config, integer_to_binary(I)) || I <- lists:seq(1,5)],
     ok.
 
+node_channel_limit(Config) ->
+    set_node_limit(Config, channel_max_per_node, 5),
+
+    VHost = <<"foobar">>,
+    User = <<"guest">>,
+    ok = rabbit_ct_broker_helpers:add_vhost(Config, VHost),
+    ok = rabbit_ct_broker_helpers:set_full_permissions(Config, User, VHost),
+    Conn1 = rabbit_ct_client_helpers:open_unmanaged_connection(Config, 0, VHost),
+    Conn2 = rabbit_ct_client_helpers:open_unmanaged_connection(Config, 0, VHost),
+
+    lists:foreach(fun(N) when (N band 1) == 1 -> {ok, _} = open_channel(Conn1);
+                     (_) -> {ok,_ } = open_channel(Conn2)
+                  end, lists:seq(1, 5)),
+
+    5 = count_channels_per_node(Config),
+    %% In total 5 channels are open on this node, so a new one, regardless of
+    %% connection, will not be allowed. It will terminate the connection with
+    %% its channels too. So
+    {error, not_allowed_crash} = open_channel(Conn2),
+    3 = count_channels_per_node(Config),
+    %% As the connection is dead, so are the 2 channels, so we should be able to
+    %% create 2 more on Conn1
+    {ok , _} = open_channel(Conn1),
+    {ok , _} = open_channel(Conn1),
+    %% But not a third
+    {error, not_allowed_crash} = open_channel(Conn1),
+
+    %% Now all connections are closed, so there should be 0 open connections
+    0 = count_channels_per_node(Config),
+    ok.
 
 %% -------------------------------------------------------------------
 %% Implementation
 %% -------------------------------------------------------------------
 
 open_connections_to_limit(Config, Limit) ->
-    set_node_limit(Config, Limit),
+    set_node_limit(Config, connection_max, Limit),
     Connections = [rabbit_ct_client_helpers:open_unmanaged_connection(Config, 0) || _ <- lists:seq(1,Limit)],
     true = lists:all(fun(E) -> is_pid(E) end, Connections),
     Connections.
 
 close_all_connections(Connections) ->
     [rabbit_ct_client_helpers:close_connection(C) || C <- Connections].
 
-set_node_limit(Config, Limit) ->
-    rabbit_ct_broker_helpers:rpc(Config, 0,
-                                 application,
-                                 set_env, [rabbit, connection_max, Limit]).
-
-set_vhost_limit(Config, Limit) ->
+set_node_limit(Config, Type, Limit) ->
     rabbit_ct_broker_helpers:rpc(Config, 0,
                                  application,
-                                 set_env, [rabbit, vhost_max, Limit]).
+                                 set_env, [rabbit, Type, Limit]).
+
+open_channel(Conn) when is_pid(Conn) ->
+    try amqp_connection:open_channel(Conn) of
+      {ok, Ch} -> {ok, Ch};
+      {error, _} ->
+            {error, not_allowed}
+    catch
+      _:_Error -> {error, not_allowed_crash}
+   end.
+
+count_channels_per_node(Config)  ->
+    NodeConfig = rabbit_ct_broker_helpers:get_node_config(Config, 0),
+    length(rabbit_ct_broker_helpers:rpc(Config, 0,
+                                        rabbit_channel_tracking,
+                                        list_on_node,
+                                        [?config(nodename, NodeConfig)])).
