Merge pull request #968 from rabbitmq/rabbitmq-server-501

Provide a way to limit max number of queues in a vhost

5 files changed, 92 insertions, 27 deletions

diff --git a/src/rabbit_connection_tracking.erl b/src/rabbit_connection_tracking.erl
index b1d34ca516..9fd3244979 100644
--- a/src/rabbit_connection_tracking.erl
+++ b/src/rabbit_connection_tracking.erl
@@ -37,7 +37,7 @@
          list/0, list/1, list_on_node/1,
          tracked_connection_from_connection_created/1,
          tracked_connection_from_connection_state/1,
-         is_over_connection_limit/1, count_connections_in/1]).
+         count_connections_in/1]).
 
 -include_lib("rabbit.hrl").
 
@@ -214,29 +214,6 @@ list_on_node(Node) ->
     catch exit:{aborted, {no_exists, _}} -> []
     end.
 
--spec is_over_connection_limit(rabbit_types:vhost()) -> {true, non_neg_integer()} | false.
-
-is_over_connection_limit(VirtualHost) ->
-    case rabbit_vhost_limit:connection_limit(VirtualHost) of
-        %% no limit configured
-        undefined                                            -> false;
-        %% with limit = 0, no connections are allowed
-        {ok, 0}                                              -> {true, 0};
-        {ok, Limit} when is_integer(Limit) andalso Limit > 0 ->
-            ConnectionCount = count_connections_in(VirtualHost),
-            case ConnectionCount >= Limit of
-                false -> false;
-                true  -> {true, Limit}
-            end;
-        %% any negative value means "no limit". Note that parameter validation
-        %% will replace negative integers with 'undefined', so this is to be
-        %% explicit and extra defensive
-        {ok, Limit} when is_integer(Limit) andalso Limit < 0 -> false;
-        %% ignore non-integer limits
-        {ok, _Limit}                                         -> false
-    end.
-
-
 -spec count_connections_in(rabbit_types:vhost()) -> non_neg_integer().
 
 count_connections_in(VirtualHost) ->
diff --git a/src/rabbit_direct.erl b/src/rabbit_direct.erl
index 858681ecfd..53b0340b8a 100644
--- a/src/rabbit_direct.erl
+++ b/src/rabbit_direct.erl
@@ -106,7 +106,7 @@ is_over_connection_limit(VHost, {Username, _Password}, Pid) ->
         none -> "";
         _    -> Username
     end,
-    try rabbit_connection_tracking:is_over_connection_limit(VHost) of
+    try rabbit_vhost_limit:is_over_connection_limit(VHost) of
         false         -> false;
         {true, Limit} ->
             rabbit_log_connection:error(
diff --git a/src/rabbit_table.erl b/src/rabbit_table.erl
index 3909096964..1bb19b23da 100644
--- a/src/rabbit_table.erl
+++ b/src/rabbit_table.erl
@@ -50,8 +50,20 @@ create() ->
                                                  Tab, TabDef1, Reason}})
                           end
                   end, definitions()),
+    ensure_secondary_indexes(),
     ok.
 
+%% Sets up secondary indexes in a blank node database.
+ensure_secondary_indexes() ->
+  ensure_secondary_index(rabbit_queue, vhost),
+  ok.
+
+ensure_secondary_index(Table, Field) ->
+  case mnesia:add_table_index(Table, Field) of
+    {atomic, ok}                          -> ok;
+    {aborted, {already_exists, Table, _}} -> ok
+  end.
+
 %% The sequence in which we delete the schema and then the other
 %% tables is important: if we delete the schema first when moving to
 %% RAM mnesia will loudly complain since it doesn't make much sense to
diff --git a/src/rabbit_upgrade_functions.erl b/src/rabbit_upgrade_functions.erl
index 30fe4fa17f..a53ad0c8f9 100644
--- a/src/rabbit_upgrade_functions.erl
+++ b/src/rabbit_upgrade_functions.erl
@@ -57,6 +57,7 @@
 -rabbit_upgrade({user_password_hashing, mnesia, [hash_passwords]}).
 -rabbit_upgrade({operator_policies,     mnesia, [slave_pids_pending_shutdown, internal_system_x]}).
 -rabbit_upgrade({vhost_limits,          mnesia, []}).
+-rabbit_upgrade({queue_vhost_field,     mnesia, [operator_policies]}).
 
 %% -------------------------------------------------------------------
 
@@ -91,6 +92,8 @@
 -spec recoverable_slaves() -> 'ok'.
 -spec user_password_hashing() -> 'ok'.
 -spec vhost_limits() -> 'ok'.
+-spec operator_policies() -> 'ok'.
+-spec queue_vhost_field() -> 'ok'.
 
 
 %%--------------------------------------------------------------------
@@ -528,6 +531,27 @@ queue_operator_policies(Table) ->
        gm_pids, decorators, state, policy_version, slave_pids_pending_shutdown]).
 
 
+queue_vhost_field() ->
+    ok = queue_vhost_field(rabbit_queue),
+    ok = queue_vhost_field(rabbit_durable_queue),
+    {atomic, ok} = mnesia:add_table_index(rabbit_queue, vhost),
+    {atomic, ok} = mnesia:add_table_index(rabbit_durable_queue, vhost),
+    ok.
+
+queue_vhost_field(Table) ->
+    transform(
+      Table,
+      fun ({amqqueue, Name = {resource, VHost, queue, _QName}, Durable, AutoDelete, ExclusiveOwner, Arguments,
+            Pid, SlavePids, SyncSlavePids, DSN, Policy, OperatorPolicy, GmPids, Decorators,
+            State, PolicyVersion, SlavePidsPendingShutdown}) ->
+              {amqqueue, Name, Durable, AutoDelete, ExclusiveOwner, Arguments,
+               Pid, SlavePids, SyncSlavePids, DSN, Policy, OperatorPolicy, GmPids, Decorators,
+               State, PolicyVersion, SlavePidsPendingShutdown, VHost}
+      end,
+      [name, durable, auto_delete, exclusive_owner, arguments, pid, slave_pids,
+       sync_slave_pids, recoverable_slaves, policy, operator_policy,
+       gm_pids, decorators, state, policy_version, slave_pids_pending_shutdown, vhost]).
+
 %% Prior to 3.6.0, passwords were hashed using MD5, this populates
 %% existing records with said default.  Users created with 3.6.0+ will
 %% have internal_user.hashing_algorithm populated by the internal
diff --git a/src/rabbit_vhost_limit.erl b/src/rabbit_vhost_limit.erl
index 2d9a2f075e..bd79f4dd45 100644
--- a/src/rabbit_vhost_limit.erl
+++ b/src/rabbit_vhost_limit.erl
@@ -23,7 +23,8 @@
 -export([register/0]).
 -export([parse_set/2, clear/1]).
 -export([validate/5, notify/4, notify_clear/3]).
--export([connection_limit/1]).
+-export([connection_limit/1, queue_limit/1,
+         is_over_queue_limit/1, is_over_connection_limit/1]).
 
 -import(rabbit_misc, [pget/2]).
 
@@ -53,6 +54,56 @@ notify_clear(VHost, <<"vhost-limits">>, <<"limits">>) ->
 connection_limit(VirtualHost) ->
     get_limit(VirtualHost, <<"max-connections">>).
 
+queue_limit(VirtualHost) ->
+    get_limit(VirtualHost, <<"max-queues">>).
+
+
+-spec is_over_connection_limit(rabbit_types:vhost()) -> {true, non_neg_integer()} | false.
+
+is_over_connection_limit(VirtualHost) ->
+    case rabbit_vhost_limit:connection_limit(VirtualHost) of
+        %% no limit configured
+        undefined                                            -> false;
+        %% with limit = 0, no connections are allowed
+        {ok, 0}                                              -> {true, 0};
+        {ok, Limit} when is_integer(Limit) andalso Limit > 0 ->
+            ConnectionCount = rabbit_connection_tracking:count_connections_in(VirtualHost),
+            case ConnectionCount >= Limit of
+                false -> false;
+                true  -> {true, Limit}
+            end;
+        %% any negative value means "no limit". Note that parameter validation
+        %% will replace negative integers with 'undefined', so this is to be
+        %% explicit and extra defensive
+        {ok, Limit} when is_integer(Limit) andalso Limit < 0 -> false;
+        %% ignore non-integer limits
+        {ok, _Limit}                                         -> false
+    end.
+
+
+-spec is_over_queue_limit(rabbit_types:vhost()) -> {true, non_neg_integer()} | false.
+
+is_over_queue_limit(VirtualHost) ->
+    case queue_limit(VirtualHost) of
+        %% no limit configured
+        undefined                                            -> false;
+        %% with limit = 0, no queues can be declared (perhaps not very
+        %% useful but consistent with the connection limit)
+        {ok, 0}                                              -> {true, 0};
+        {ok, Limit} when is_integer(Limit) andalso Limit > 0 ->
+            QueueCount = rabbit_amqqueue:count(VirtualHost),
+            case QueueCount >= Limit of
+                false -> false;
+                true  -> {true, Limit}
+            end;
+        %% any negative value means "no limit". Note that parameter validation
+        %% will replace negative integers with 'undefined', so this is to be
+        %% explicit and extra defensive
+        {ok, Limit} when is_integer(Limit) andalso Limit < 0 -> false;
+        %% ignore non-integer limits
+        {ok, _Limit}                                         -> false
+    end.
+
 %%----------------------------------------------------------------------------
 
 parse_set(VHost, Defn) ->
@@ -72,7 +123,8 @@ clear(VHost) ->
                                         <<"limits">>).
 
 vhost_limit_validation() ->
-    [{<<"max-connections">>, fun rabbit_parameter_validation:integer/2, mandatory}].
+    [{<<"max-connections">>, fun rabbit_parameter_validation:integer/2, optional},
+     {<<"max-queues">>,      fun rabbit_parameter_validation:integer/2, optional}].
 
 update_vhost(VHostName, Limits) ->
     rabbit_misc:execute_mnesia_transaction(