Support setting vm_memory_high_watermark and disk_free_limit in memory units

4 files changed, 67 insertions, 8 deletions

diff --git a/src/rabbit_alarm.erl b/src/rabbit_alarm.erl
index 111f780076..09d67516fd 100644
--- a/src/rabbit_alarm.erl
+++ b/src/rabbit_alarm.erl
@@ -39,6 +39,11 @@
 
 -define(SERVER, ?MODULE).
 
+-define(DEFAULT_VM_MEMORY_HIGH_WATERMARK, 0.4).
+-define(DEFAULT_DISK_FREE_LIMIT, 50000000).
+
+-export([parse_limit/1]).
+
 %%----------------------------------------------------------------------------
 
 -ifdef(use_specs).
@@ -61,6 +66,8 @@
 -spec(on_node_up/1 :: (node()) -> 'ok').
 -spec(on_node_down/1 :: (node()) -> 'ok').
 -spec(get_alarms/0 :: () -> [{alarm(), []}]).
+-spec(parse_limit/1 :: (integer() | string()) -> 
+    {ok, integer()} | {error, parse_error}).
 
 -else.
 
@@ -77,8 +84,9 @@ start() ->
     ok = rabbit_sup:start_restartable_child(?MODULE),
     ok = gen_event:add_handler(?SERVER, ?MODULE, []),
     {ok, MemoryWatermark} = application:get_env(vm_memory_high_watermark),
+
     rabbit_sup:start_restartable_child(
-      vm_memory_monitor, [MemoryWatermark,
+      vm_memory_monitor, [parse_mem_limit(MemoryWatermark),
                           fun (Alarm) ->
                                   background_gc:run(),
                                   set_alarm(Alarm)
@@ -86,9 +94,56 @@ start() ->
                           fun clear_alarm/1]),
     {ok, DiskLimit} = application:get_env(disk_free_limit),
     rabbit_sup:start_delayed_restartable_child(
-      rabbit_disk_monitor, [DiskLimit]),
+      rabbit_disk_monitor, [parse_disk_limit(DiskLimit)]),
     ok.
 
+parse_mem_limit({absolute, Limit}) ->
+    case parse_limit(Limit) of
+        {ok, ParsedLimit} -> {absolute, ParsedLimit};
+        {error, parse_error} ->
+            rabbit_log:error("Unable to parse vm_memory_high_watermark value"),
+            ?DEFAULT_VM_MEMORY_HIGH_WATERMARK
+    end;
+    {absolute, parse_limit(Limit, ?DEFAULT_VM_MEMORY_HIGH_WATERMARK)};
+parse_mem_limit(Relative) when is_float(Relative), Relative < 1 ->
+    Relative;
+parse_mem_limit(_) ->
+    ?DEFAULT_VM_MEMORY_HIGH_WATERMARK.
+
+parse_disk_limit({mem_relative, Relative}) 
+    when is_float(Relative), Relative < 1 ->
+    {mem_relative, Relative};
+parse_disk_limit(Limit) -> 
+    case parse_limit(Limit) of
+        {ok, ParsedLimit} -> ParsedLimit;
+        {error, parse_error} ->
+            rabbit_log:error("Unable to parse disk_free_limit value"),
+            ?DEFAULT_DISK_FREE_LIMIT
+    end;
+parse_disk_limit(_) -> ?DEFAULT_DISK_FREE_LIMIT.
+
+parse_limit(MemLim) when is_integer(MemLim) -> {ok, MemLim};
+parse_limit(MemLim) when is_string(MemLim) ->
+    case re:run(MemLim, 
+                "^(?<VAL>[1-9][0-9]*)(?<UNIT>kB|MB|GB|kiB|MiB|GiB|k|M|G)$", 
+                [{capture, all_names}]) of
+        {match, [{NumPos, NumLength}, {UnitPos, UnitLength}]} ->
+            Num = list_to_integer(string:substr(MemLim, NumPos+1, NumLength))
+            Unit = string:substr(MemLim, UnitPos+1, UniqLength),
+            Multiplier = case Unit of
+                KiB when KiB == "k"; KiB == "kiB" -> 1024;
+                MiB when MiB == "M"; MiB == "MiB" -> 1024*1024;
+                GiB when GiB == "G"; GiB == "GiB" -> 1024*1024*1024;
+                "KB" -> 1000;
+                "MB" -> 1000000;
+                "GB" -> 1000000000
+            end,
+            {ok, Num * Multiplier};
+        nomatch ->
+            % log error
+            {error, parse_error}
+    end.
+
 stop() -> ok.
 
 %% Registers a handler that should be called on every resource alarm change.
diff --git a/src/rabbit_control_main.erl b/src/rabbit_control_main.erl
index d21f60a440..426520c31d 100644
--- a/src/rabbit_control_main.erl
+++ b/src/rabbit_control_main.erl
@@ -417,10 +417,14 @@ action(set_vm_memory_high_watermark, Node, [Arg], _Opts, Inform) ->
     rpc_call(Node, vm_memory_monitor, set_vm_memory_high_watermark, [Frac]);
 
 action(set_vm_memory_high_watermark, Node, ["absolute", Arg], _Opts, Inform) ->
-    Limit = list_to_integer(Arg),
-    Inform("Setting memory threshold on ~p to ~p bytes", [Node, Limit]),
-    rpc_call(Node, vm_memory_monitor, set_vm_memory_high_watermark,
-	     [{absolute, Limit}]);
+    case rabbit_alarm:parse_limit(Arg) of
+        {ok, Limit} ->
+            Inform("Setting memory threshold on ~p to ~p bytes", [Node, Limit]),
+            rpc_call(Node, vm_memory_monitor, set_vm_memory_high_watermark,
+                 [{absolute, Limit}]);
+        {error, parse_error} ->
+            {error_string, "Unable to parse absolute memory value"}
+    end;
 
 action(set_permissions, Node, [Username, CPerm, WPerm, RPerm], Opts, Inform) ->
     VHost = proplists:get_value(?VHOST_OPT, Opts),
diff --git a/src/rabbit_disk_monitor.erl b/src/rabbit_disk_monitor.erl
index 3e9171b79a..88b7469cc0 100644
--- a/src/rabbit_disk_monitor.erl
+++ b/src/rabbit_disk_monitor.erl
@@ -71,7 +71,7 @@
 
 -ifdef(use_specs).
 
--type(disk_free_limit() :: (integer() | {'mem_relative', float()})).
+-type(disk_free_limit() :: (integer() | string() | {'mem_relative', float()})).
 -spec(start_link/1 :: (disk_free_limit()) -> rabbit_types:ok_pid_or_error()).
 -spec(get_disk_free_limit/0 :: () -> integer()).
 -spec(set_disk_free_limit/1 :: (disk_free_limit()) -> 'ok').
diff --git a/src/vm_memory_monitor.erl b/src/vm_memory_monitor.erl
index 1d98843365..f07ccce89f 100644
--- a/src/vm_memory_monitor.erl
+++ b/src/vm_memory_monitor.erl
@@ -63,7 +63,7 @@
 
 -ifdef(use_specs).
 
--type(vm_memory_high_watermark() :: (float() | {'absolute', integer()})).
+-type(vm_memory_high_watermark() :: (float() | {'absolute', integer() | string()})).
 -spec(start_link/1 :: (float()) -> rabbit_types:ok_pid_or_error()).
 -spec(start_link/3 :: (float(), fun ((any()) -> 'ok'),
                        fun ((any()) -> 'ok')) -> rabbit_types:ok_pid_or_error()).