Merge pull request #3895 from rabbitmq/lukebakken/rabbit_disk_monitor-improvements

Disk monitor improvements

3 files changed, 154 insertions, 56 deletions

diff --git a/deps/rabbit/src/rabbit_disk_monitor.erl b/deps/rabbit/src/rabbit_disk_monitor.erl
index 76edecf5d8..28b4cb3eba 100644
--- a/deps/rabbit/src/rabbit_disk_monitor.erl
+++ b/deps/rabbit/src/rabbit_disk_monitor.erl
@@ -33,6 +33,7 @@
          get_disk_free/0, set_enabled/1]).
 
 -define(SERVER, ?MODULE).
+-define(ETS_NAME, ?MODULE).
 -define(DEFAULT_MIN_DISK_CHECK_INTERVAL, 100).
 -define(DEFAULT_MAX_DISK_CHECK_INTERVAL, 10000).
 -define(DEFAULT_DISK_FREE_LIMIT, 50000000).
@@ -73,51 +74,42 @@
 %%----------------------------------------------------------------------------
 
 -spec get_disk_free_limit() -> integer().
-
 get_disk_free_limit() ->
-    gen_server:call(?MODULE, get_disk_free_limit, infinity).
+    safe_ets_lookup(disk_free_limit, ?DEFAULT_DISK_FREE_LIMIT).
 
 -spec set_disk_free_limit(disk_free_limit()) -> 'ok'.
-
 set_disk_free_limit(Limit) ->
-    gen_server:call(?MODULE, {set_disk_free_limit, Limit}, infinity).
+    gen_server:call(?MODULE, {set_disk_free_limit, Limit}).
 
 -spec get_min_check_interval() -> integer().
-
 get_min_check_interval() ->
-    gen_server:call(?MODULE, get_min_check_interval, infinity).
+    safe_ets_lookup(min_check_interval, ?DEFAULT_MIN_DISK_CHECK_INTERVAL).
 
 -spec set_min_check_interval(integer()) -> 'ok'.
-
 set_min_check_interval(Interval) ->
-    gen_server:call(?MODULE, {set_min_check_interval, Interval}, infinity).
+    gen_server:call(?MODULE, {set_min_check_interval, Interval}).
 
 -spec get_max_check_interval() -> integer().
-
 get_max_check_interval() ->
-    gen_server:call(?MODULE, get_max_check_interval, infinity).
+    safe_ets_lookup(max_check_interval, ?DEFAULT_MAX_DISK_CHECK_INTERVAL).
 
 -spec set_max_check_interval(integer()) -> 'ok'.
-
 set_max_check_interval(Interval) ->
-    gen_server:call(?MODULE, {set_max_check_interval, Interval}, infinity).
+    gen_server:call(?MODULE, {set_max_check_interval, Interval}).
 
 -spec get_disk_free() -> (integer() | 'unknown').
-
 get_disk_free() ->
-    gen_server:call(?MODULE, get_disk_free, infinity).
+    safe_ets_lookup(disk_free, unknown).
 
 -spec set_enabled(string()) -> 'ok'.
-
 set_enabled(Enabled) ->
-    gen_server:call(?MODULE, {set_enabled, Enabled}, infinity).
+    gen_server:call(?MODULE, {set_enabled, Enabled}).
 
 %%----------------------------------------------------------------------------
 %% gen_server callbacks
 %%----------------------------------------------------------------------------
 
 -spec start_link(disk_free_limit()) -> rabbit_types:ok_pid_or_error().
-
 start_link(Args) ->
     gen_server:start_link({local, ?SERVER}, ?MODULE, [Args], []).
 
@@ -125,18 +117,16 @@ init([Limit]) ->
     Dir = dir(),
     {ok, Retries} = application:get_env(rabbit, disk_monitor_failure_retries),
     {ok, Interval} = application:get_env(rabbit, disk_monitor_failure_retry_interval),
-    State = #state{dir          = Dir,
-                   min_interval = ?DEFAULT_MIN_DISK_CHECK_INTERVAL,
-                   max_interval = ?DEFAULT_MAX_DISK_CHECK_INTERVAL,
-                   alarmed      = false,
-                   enabled      = true,
-                   limit        = Limit,
-                   retries      = Retries,
-                   interval     = Interval},
-    {ok, enable(State)}.
-
-handle_call(get_disk_free_limit, _From, State = #state{limit = Limit}) ->
-    {reply, Limit, State};
+    ?ETS_NAME = ets:new(?ETS_NAME, [protected, set, named_table]),
+    State0 = #state{dir          = Dir,
+                    alarmed      = false,
+                    enabled      = true,
+                    limit        = Limit,
+                    retries      = Retries,
+                    interval     = Interval},
+    State1 = set_min_check_interval(?DEFAULT_MIN_DISK_CHECK_INTERVAL, State0),
+    State2 = set_max_check_interval(?DEFAULT_MAX_DISK_CHECK_INTERVAL, State1),
+    {ok, enable(State2)}.
 
 handle_call({set_disk_free_limit, _}, _From, #state{enabled = false} = State) ->
     rabbit_log:info("Cannot set disk free limit: "
@@ -146,20 +136,14 @@ 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(get_min_check_interval, _From, State) ->
-    {reply, State#state.min_interval, State};
-
 handle_call(get_max_check_interval, _From, State) ->
     {reply, State#state.max_interval, State};
 
 handle_call({set_min_check_interval, MinInterval}, _From, State) ->
-    {reply, ok, State#state{min_interval = MinInterval}};
+    {reply, ok, set_min_check_interval(MinInterval, State)};
 
 handle_call({set_max_check_interval, MaxInterval}, _From, State) ->
-    {reply, ok, State#state{max_interval = MaxInterval}};
-
-handle_call(get_disk_free, _From, State = #state { actual = Actual }) ->
-    {reply, Actual, State};
+    {reply, ok, set_max_check_interval(MaxInterval, State)};
 
 handle_call({set_enabled, _Enabled = true}, _From, State) ->
     start_timer(set_disk_limits(State, State#state.limit)),
@@ -194,14 +178,36 @@ code_change(_OldVsn, State, _Extra) ->
 %% Server Internals
 %%----------------------------------------------------------------------------
 
+safe_ets_lookup(Key, Default) ->
+    try
+        case ets:lookup(?ETS_NAME, Key) of
+            [{Key, Value}] ->
+                Value;
+            [] ->
+                Default
+        end
+    catch
+        error:badarg ->
+            Default
+    end.
+
 % the partition / drive containing this directory will be monitored
 dir() -> rabbit_mnesia:dir().
 
+set_min_check_interval(MinInterval, State) ->
+    ets:insert(?ETS_NAME, {min_check_interval, MinInterval}),
+    State#state{min_interval = MinInterval}.
+
+set_max_check_interval(MaxInterval, State) ->
+    ets:insert(?ETS_NAME, {max_check_interval, MaxInterval}),
+    State#state{max_interval = MaxInterval}.
+
 set_disk_limits(State, Limit0) ->
     Limit = interpret_limit(Limit0),
     State1 = State#state { limit = Limit },
     rabbit_log:info("Disk free limit set to ~pMB",
                     [trunc(Limit / 1000000)]),
+    ets:insert(?ETS_NAME, {disk_free_limit, Limit}),
     internal_update(State1).
 
 internal_update(State = #state { limit   = Limit,
@@ -219,7 +225,8 @@ internal_update(State = #state { limit   = Limit,
         _ ->
             ok
     end,
-    State #state {alarmed = NewAlarmed, actual = CurrentFree}.
+    ets:insert(?ETS_NAME, {disk_free, CurrentFree}),
+    State#state{alarmed = NewAlarmed, actual = CurrentFree}.
 
 get_disk_free(Dir) ->
     get_disk_free(Dir, os:type()).
@@ -227,11 +234,89 @@ get_disk_free(Dir) ->
 get_disk_free(Dir, {unix, Sun})
   when Sun =:= sunos; Sun =:= sunos4; Sun =:= solaris ->
     Df = os:find_executable("df"),
-    parse_free_unix(rabbit_misc:os_cmd(Df ++ " -k " ++ Dir));
+    parse_free_unix(run_cmd(Df ++ " -k " ++ Dir));
 get_disk_free(Dir, {unix, _}) ->
     Df = os:find_executable("df"),
-    parse_free_unix(rabbit_misc:os_cmd(Df ++ " -kP " ++ Dir));
+    parse_free_unix(run_cmd(Df ++ " -kP " ++ Dir));
 get_disk_free(Dir, {win32, _}) ->
+    % Dir:
+    % "c:/Users/username/AppData/Roaming/RabbitMQ/db/rabbit2@username-z01-mnesia"
+    case win32_get_drive_letter(Dir) of
+        error ->
+            rabbit_log:warning("Expected the mnesia directory absolute "
+                               "path to start with a drive letter like "
+                               "'C:'. The path is: '~p'", [Dir]),
+            case win32_get_disk_free_dir(Dir) of
+                {ok, Free} ->
+                    Free;
+                _ -> exit(could_not_determine_disk_free)
+            end;
+        DriveLetter ->
+            case win32_get_disk_free_fsutil(DriveLetter) of
+                {ok, Free0} -> Free0;
+                error ->
+                    case win32_get_disk_free_pwsh(DriveLetter) of
+                        {ok, Free1} -> Free1;
+                        _ -> exit(could_not_determine_disk_free)
+                    end
+            end
+    end.
+
+parse_free_unix(Str) ->
+    case string:tokens(Str, "\n") of
+        [_, S | _] -> case string:tokens(S, " \t") of
+                          [_, _, _, Free | _] -> list_to_integer(Free) * 1024;
+                          _                   -> exit({unparseable, Str})
+                      end;
+        _          -> exit({unparseable, Str})
+    end.
+
+win32_get_drive_letter([DriveLetter, $:, $/ | _]) when
+      (DriveLetter >= $a andalso DriveLetter =< $z) orelse
+      (DriveLetter >= $A andalso DriveLetter =< $Z) ->
+    DriveLetter;
+win32_get_drive_letter(_) ->
+    error.
+
+win32_get_disk_free_fsutil(DriveLetter) when
+      (DriveLetter >= $a andalso DriveLetter =< $z) orelse
+      (DriveLetter >= $A andalso DriveLetter =< $Z) ->
+    % DriveLetter $c
+    FsutilCmd = "fsutil.exe volume diskfree " ++ [DriveLetter] ++ ":",
+
+    % C:\windows\system32>fsutil volume diskfree c:
+    % Total free bytes        :   812,733,878,272 (756.9 GB)
+    % Total bytes             : 1,013,310,287,872 (943.7 GB)
+    % Total quota free bytes  :   812,733,878,272 (756.9 GB)
+    case run_cmd(FsutilCmd) of
+        {error, timeout} ->
+            error;
+        FsutilResult ->
+            case string:slice(FsutilResult, 0, 5) of
+                "Error" ->
+                    error;
+                "Total" ->
+                    FirstLine = hd(string:tokens(FsutilResult, "\r\n")),
+                    {match, [FreeStr]} = re:run(FirstLine, "(\\d+,?)+", [{capture, first, list}]),
+                    {ok, list_to_integer(lists:flatten(string:tokens(FreeStr, ",")))}
+            end
+    end.
+
+win32_get_disk_free_pwsh(DriveLetter) when
+      (DriveLetter >= $a andalso DriveLetter =< $z) orelse
+      (DriveLetter >= $A andalso DriveLetter =< $Z) ->
+    % DriveLetter $c
+    PoshCmd = "powershell.exe -NoLogo -NoProfile -NonInteractive -Command (Get-PSDrive " ++ [DriveLetter] ++ ").Free",
+    case run_cmd(PoshCmd) of
+        {error, timeout} ->
+            error;
+        PoshResultStr ->
+            % Note: remove \r\n
+            PoshResult = string:slice(PoshResultStr, 0, length(PoshResultStr) - 2),
+            {ok, list_to_integer(PoshResult)}
+    end.
+
+win32_get_disk_free_dir(Dir) ->
     %% On Windows, the Win32 API enforces a limit of 260 characters
     %% (MAX_PATH). If we call `dir` with a path longer than that, it
     %% fails with "File not found". Starting with Windows 10 version
@@ -253,22 +338,11 @@ get_disk_free(Dir, {win32, _}) ->
     %% See the following page to learn more about this:
     %% https://ss64.com/nt/syntax-filenames.html
     RawDir = "\\\\?\\" ++ string:replace(Dir, "/", "\\", all),
-    parse_free_win32(rabbit_misc:os_cmd("dir /-C /W \"" ++ RawDir ++ "\"")).
-
-parse_free_unix(Str) ->
-    case string:tokens(Str, "\n") of
-        [_, S | _] -> case string:tokens(S, " \t") of
-                          [_, _, _, Free | _] -> list_to_integer(Free) * 1024;
-                          _                   -> exit({unparseable, Str})
-                      end;
-        _          -> exit({unparseable, Str})
-    end.
-
-parse_free_win32(CommandResult) ->
+    CommandResult = run_cmd("dir /-C /W \"" ++ RawDir ++ "\""),
     LastLine = lists:last(string:tokens(CommandResult, "\r\n")),
     {match, [Free]} = re:run(lists:reverse(LastLine), "(\\d+)",
                              [{capture, all_but_first, list}]),
-    list_to_integer(lists:reverse(Free)).
+    {ok, list_to_integer(lists:reverse(Free))}.
 
 interpret_limit({mem_relative, Relative})
     when is_number(Relative) ->
@@ -318,3 +392,20 @@ enable(#state{dir = Dir, interval = Interval, limit = Limit, retries = Retries}
             erlang:send_after(Interval, self(), try_enable),
             State#state{enabled = false}
     end.
+
+run_cmd(Cmd) ->
+    Pid = self(),
+    Ref = make_ref(),
+    CmdFun = fun() ->
+        CmdResult = rabbit_misc:os_cmd(Cmd),
+        Pid ! {Pid, Ref, CmdResult}
+    end,
+    CmdPid = spawn(CmdFun),
+    receive
+        {Pid, Ref, CmdResult} ->
+            CmdResult
+    after 5000 ->
+        exit(CmdPid, kill),
+        rabbit_log:error("Command timed out: '~s'", [Cmd]),
+        {error, timeout}
+    end.
diff --git a/deps/rabbit/test/unit_disk_monitor_SUITE.erl b/deps/rabbit/test/unit_disk_monitor_SUITE.erl
index 71393d5d35..7a746349ca 100644
--- a/deps/rabbit/test/unit_disk_monitor_SUITE.erl
+++ b/deps/rabbit/test/unit_disk_monitor_SUITE.erl
@@ -67,6 +67,12 @@ set_disk_free_limit_command(Config) ->
       ?MODULE, set_disk_free_limit_command1, [Config]).
 
 set_disk_free_limit_command1(_Config) ->
+    F = fun () ->
+        DiskFree = rabbit_disk_monitor:get_disk_free(),
+        DiskFree =/= unknown
+    end,
+    rabbit_ct_helpers:await_condition(F),
+
     %% Use an integer
     rabbit_disk_monitor:set_disk_free_limit({mem_relative, 1}),
     disk_free_limit_to_total_memory_ratio_is(1),
@@ -84,7 +90,8 @@ set_disk_free_limit_command1(_Config) ->
     passed.
 
 disk_free_limit_to_total_memory_ratio_is(MemRatio) ->
+    DiskFreeLimit = rabbit_disk_monitor:get_disk_free_limit(),
     ExpectedLimit = MemRatio * vm_memory_monitor:get_total_memory(),
     % Total memory is unstable, so checking order
-    true = ExpectedLimit/rabbit_disk_monitor:get_disk_free_limit() < 1.2,
-    true = ExpectedLimit/rabbit_disk_monitor:get_disk_free_limit() > 0.98.
+    true = ExpectedLimit/DiskFreeLimit < 1.2,
+    true = ExpectedLimit/DiskFreeLimit > 0.98.
diff --git a/deps/rabbit/test/unit_disk_monitor_mocks_SUITE.erl b/deps/rabbit/test/unit_disk_monitor_mocks_SUITE.erl
index 80fa3d2e09..ae16cbb379 100644
--- a/deps/rabbit/test/unit_disk_monitor_mocks_SUITE.erl
+++ b/deps/rabbit/test/unit_disk_monitor_mocks_SUITE.erl
@@ -88,7 +88,7 @@ disk_monitor_enable1(_Config) ->
     application:set_env(rabbit, disk_monitor_failure_retry_interval, 100),
     ok = rabbit_sup:stop_child(rabbit_disk_monitor_sup),
     ok = rabbit_sup:start_delayed_restartable_child(rabbit_disk_monitor, [1000]),
-    undefined = rabbit_disk_monitor:get_disk_free(),
+    unknown = rabbit_disk_monitor:get_disk_free(),
     Cmd = case os:type() of
               {win32, _} -> " Le volume dans le lecteur C n’a pas de nom.\n"
                             " Le numéro de série du volume est 707D-5BDC\n"