Enforce max connection limit

Also introduce `rabbitmqctl clear_vhost_limits`
and fix rabbitmqctl(1).

docs/rabbitmqctl.1.xml

@@ -1194,27 +1194,45 @@
       </para>
       <variablelist>
         <varlistentry>
-          <term><cmdsynopsis><command>set_vhost_limits</command> <arg choice="opt">-p <replaceable>vhostpath</replaceable></arg> <arg choice="opt">--max-connections <replaceable>number</replaceable></arg></cmdsynopsis></term>
+          <term><cmdsynopsis><command>set_vhost_limits</command> <arg choice="opt">-p <replaceable>vhostpath</replaceable></arg> <arg choice="req"><replaceable>definition</replaceable></arg></cmdsynopsis></term>
           <listitem>
             <para>
               Sets virtual host limits
             </para>
             <variablelist>
               <varlistentry>
-                <term>max-connections</term>
+                <term>definition</term>
                 <listitem><para>
-                    How many concurrent client connections are allowed. 0 means "no limit". Default is "no limit".
+                    The definition of the limits, as a
+                    JSON term. In most shells you are very likely to
+                    need to quote this.
+
+                    Recognised limits: max-connections (0 means "no limit").
                 </para></listitem>
               </varlistentry>
             </variablelist>
             <para role="example-prefix">For example:</para>
-            <screen role="example">rabbitmqctl set_vhost_limits -p qa_env --max-connections 1024</screen>
+            <screen role="example">rabbitmqctl set_vhost_limits -p qa_env '{"max-connections": 1024}'</screen>
             <para role="example">
               This command limits the max number of concurrent connections in vhost <command>qa_env</command>
               to 1024.
             </para>
           </listitem>
         </varlistentry>
+
+<varlistentry>
+          <term><cmdsynopsis><command>clear_vhost_limits</command> <arg choice="opt">-p <replaceable>vhostpath</replaceable></arg></cmdsynopsis></term>
+          <listitem>
+            <para>
+              Clears virtual host limits
+            </para>
+            <para role="example-prefix">For example:</para>
+            <screen role="example">rabbitmqctl clear_vhost_limits -p qa_env</screen>
+            <para role="example">
+              This command clears vhost limits in vhost <command>qa_env</command>.
+            </para>
+          </listitem>
+        </varlistentry>
       </variablelist>
     </refsect2>
 

src/rabbit_connection_tracking.erl

@@ -26,7 +26,8 @@
 %%  * rabbit_event
 
 -export([register_connection/1, unregister_connection/1,
-         tracked_connection_from_connection_created/1]).
+         tracked_connection_from_connection_created/1,
+         is_over_connection_limit/1, count_connections_in/1]).
 
 -ifdef(use_specs).
 
@@ -54,6 +55,26 @@ unregister_connection(ConnId = {_Node, _Name}) ->
         mnesia:delete({?TABLE, ConnId})
     end).
 
+is_over_connection_limit(VirtualHost) ->
+    ConnectionCount = count_connections_in(VirtualHost),
+    case rabbit_vhost_limit:connection_limit(VirtualHost) of
+        undefined   -> false;
+        {ok, Limit} -> case ConnectionCount > Limit of
+                           false -> false;
+                           true  -> {true, Limit}
+                       end
+    end.
+
+count_connections_in(VirtualHost) ->
+    %% TODO: optimize
+    Xs = rabbit_misc:execute_mnesia_transaction(
+        fun() ->
+            mnesia:index_read(
+                rabbit_tracked_connection, VirtualHost,
+                #tracked_connection.vhost)
+        end),
+    length(Xs).
+
 %% Returns a #tracked_connection from connection_created
 %% event details.
 %%

src/rabbit_control_main.erl

@@ -73,6 +73,7 @@
          {list_policies, [?VHOST_DEF]},
 
          {set_vhost_limits, [?VHOST_DEF]},
+         {clear_vhost_limits, [?VHOST_DEF]},
 
          {list_queues, [?VHOST_DEF]},
          {list_exchanges, [?VHOST_DEF]},
@@ -521,6 +522,11 @@ action(set_vhost_limits, Node, [Defn], Opts, Inform) ->
     rpc_call(Node, rabbit_vhost_limit, parse_set, [VHostArg, Defn]),
     ok;
 
+action(clear_vhost_limits, Node, [], Opts, Inform) ->
+    VHostArg = list_to_binary(proplists:get_value(?VHOST_OPT, Opts)),
+    Inform("Clearing vhost ~p limits", [VHostArg]),
+    rpc_call(Node, rabbit_vhost_limit, clear, [VHostArg]);
+
 action(report, Node, _Args, _Opts, Inform) ->
     Inform("Reporting server status on ~p~n~n", [erlang:universaltime()]),
     [begin ok = action(Action, N, [], [], Inform), io:nl() end ||

src/rabbit_vhost_limit.erl

@@ -21,8 +21,11 @@
 -include("rabbit.hrl").
 
 -export([register/0]).
--export([parse_set/2]).
+-export([parse_set/2, clear/1]).
 -export([validate/5, notify/4, notify_clear/3]).
+-export([connection_limit/1]).
+
+-import(rabbit_misc, [pget/2]).
 
 -rabbit_boot_step({?MODULE,
                    [{description, "vhost limit parameters"},
@@ -47,6 +50,9 @@ notify_clear(VHost, <<"vhost-limits">>, <<"limits">>) ->
     rabbit_event:notify(vhost_limits_cleared, [{name, <<"limits">>}]),
     update_vhost(VHost, undefined).
 
+connection_limit(VirtualHost) ->
+    get_limit(VirtualHost, <<"max-connections">>).
+
 %%----------------------------------------------------------------------------
 
 parse_set(VHost, Defn) ->
@@ -61,6 +67,10 @@ set(VHost, Defn) ->
     rabbit_runtime_parameters:set_any(VHost, <<"vhost-limits">>,
                                       <<"limits">>, Defn, none).
 
+clear(VHost) ->
+    rabbit_runtime_parameters:clear_any(VHost, <<"vhost-limits">>,
+                                        <<"limits">>).
+
 vhost_limit_validation() ->
     [{<<"max-connections">>, fun rabbit_parameter_validation:number/2, mandatory}].
 
@@ -73,3 +83,16 @@ update_vhost(VHostName, Limits) ->
                                   end)
       end),
     ok.
+
+get_limit(VirtualHost, Limit) ->
+    case rabbit_runtime_parameters:list(VirtualHost, <<"vhost-limits">>) of
+        []      -> undefined;
+        [Param] -> case pget(value, Param) of
+                       undefined -> undefined;
+                       Val       -> case pget(Limit, Val) of
+                                        undefined     -> undefined;
+                                        N when N =< 0 -> undefined;
+                                        N when N > 0  -> {ok, N}
+                                    end
+                   end
+    end.