CLI: Extend set_disk_free_limit to set mount limits

With this change you can say:

    rabbitmqctl set_disk_limit mount Streaming 2GiB

This applies the limit only to the "Streaming" mount.

Author: the-mikedavis

deps/rabbit/src/rabbit_disk_monitor.erl

@@ -26,6 +26,7 @@
          terminate/2, code_change/3]).
 
 -export([get_disk_free_limit/0, set_disk_free_limit/1,
+         set_disk_free_limit/2,
          get_min_check_interval/0, set_min_check_interval/1,
          get_max_check_interval/0, set_max_check_interval/1,
          get_disk_free/0, get_mount_free/0, set_enabled/1]).
@@ -105,6 +106,11 @@ get_disk_free_limit() ->
 set_disk_free_limit(Limit) ->
     gen_server:call(?MODULE, {set_disk_free_limit, Limit}).
 
+-spec set_disk_free_limit(MountName :: binary(), integer()) -> 'ok'.
+set_disk_free_limit(MountName, Limit)
+  when is_binary(MountName) andalso is_integer(Limit) ->
+    gen_server:call(?MODULE, {set_disk_free_limit, MountName, Limit}).
+
 -spec get_min_check_interval() -> integer().
 get_min_check_interval() ->
     safe_ets_lookup(min_check_interval, ?DEFAULT_MIN_DISK_CHECK_INTERVAL).
@@ -189,6 +195,23 @@ handle_call({set_disk_free_limit, _}, _From, #state{enabled = false} = State) ->
 handle_call({set_disk_free_limit, Limit}, _From, State) ->
     {reply, ok, set_disk_limits(State, Limit)};
 
+handle_call({set_disk_free_limit, Name, Limit}, _From,
+            #state{mounts = Mounts0} = State) ->
+    MatchingMount = lists:search(
+                      fun({_Path, #mount{name = N}}) ->
+                              Name =:= N
+                      end, maps:to_list(Mounts0)),
+    case MatchingMount of
+        {value, {Path, Mount}} ->
+            ?LOG_INFO("Updated disk free limit of mount '~ts'", [Name]),
+            Mounts = Mounts0#{Path := Mount#mount{limit = Limit}},
+            {reply, ok, State#state{mounts = Mounts}};
+        false ->
+            ?LOG_WARNING("Cannot set disk free limit for mount '~ts' since "
+                         "the name does not match any known mounts.", [Name]),
+            {reply, ok, State}
+    end;
+
 handle_call(get_max_check_interval, _From, State) ->
     {reply, State#state.max_interval, State};
 

deps/rabbitmq_cli/lib/rabbitmq/cli/ctl/commands/set_disk_free_limit_command.ex

@@ -24,6 +24,14 @@ defmodule RabbitMQ.CLI.Ctl.Commands.SetDiskFreeLimitCommand do
     {:validation_failure, :too_many_args}
   end
 
+  def validate(["mount" | _] = args, _) when length(args) < 3 do
+    {:validation_failure, :not_enough_args}
+  end
+
+  def validate(["mount" | _] = args, _) when length(args) > 3 do
+    {:validation_failure, :too_many_args}
+  end
+
   def validate([limit], _) do
     case Integer.parse(limit) do
       {_, ""} ->
@@ -47,12 +55,41 @@ defmodule RabbitMQ.CLI.Ctl.Commands.SetDiskFreeLimitCommand do
     end
   end
 
+  def validate(["mount", _name, limit], _) do
+    case Integer.parse(limit) do
+      {_, ""} ->
+        :ok
+
+      {limit_val, units} ->
+        case memory_unit_absolute(limit_val, units) do
+          scaled_limit when is_integer(scaled_limit) -> :ok
+          _ -> {:validation_failure, :bad_argument}
+        end
+
+      _ ->
+        {:validation_failure, :bad_argument}
+    end
+  end
+
   def validate([_ | rest], _) when length(rest) > 0 do
     {:validation_failure, :too_many_args}
   end
 
   use RabbitMQ.CLI.Core.RequiresRabbitAppRunning
 
+  def run(["mount", mount_name, limit], %{node: node_name}) do
+    limit =
+      case Integer.parse(limit) do
+        {limit, ""} -> limit
+        {limit, units} ->
+          case memory_unit_absolute(limit, units) do
+            scaled_limit when is_integer(scaled_limit) ->
+              scaled_limit
+          end
+      end
+    make_rpc_call(node_name, [mount_name, limit])
+  end
+
   def run(["mem_relative", _] = args, opts) do
     set_disk_free_limit_relative(args, opts)
   end
@@ -70,19 +107,24 @@ defmodule RabbitMQ.CLI.Ctl.Commands.SetDiskFreeLimitCommand do
 
   use RabbitMQ.CLI.DefaultOutput
 
+  def banner(["mount", mount, limit], %{node: node_name}) do
+    "Setting disk free limit for mount #{mount} on #{node_name} to #{limit} ..."
+  end
+
   def banner(["mem_relative", arg], %{node: node_name}) do
     "Setting disk free limit on #{node_name} to #{arg} times the total RAM ..."
   end
 
   def banner([arg], %{node: node_name}),
     do: "Setting disk free limit on #{node_name} to #{arg} bytes ..."
 
-  def usage, do: "set_disk_free_limit <disk_limit> | mem_relative <fraction>"
+  def usage, do: "set_disk_free_limit <disk_limit> | mem_relative <fraction> | mount <name> <disk_limit>"
 
   def usage_additional() do
     [
       ["<disk_limit>", "New limit as an absolute value with units, e.g. 1GB"],
-      ["mem_relative <fraction>", "New limit as a fraction of total memory reported by the OS"]
+      ["mem_relative <fraction>", "New limit as a fraction of total memory reported by the OS"],
+      ["mount <name> <disk_limit>", "New limit for the given mount name as an absolute value with units, e.g. 1GB"]
     ]
   end