Keep backup between mnesia upgrade and message store upgrade

2 files changed, 50 insertions, 12 deletions

diff --git a/src/rabbit_upgrade.erl b/src/rabbit_upgrade.erl
index 8b7d11e4fa..1ee7423404 100644
--- a/src/rabbit_upgrade.erl
+++ b/src/rabbit_upgrade.erl
@@ -108,6 +108,7 @@ ensure_backup_taken() ->
 
 take_backup() ->
     BackupDir = backup_dir(),
+    info("upgrades: Backing up mnesia dir to ~p~n", [BackupDir]),
     case rabbit_mnesia:copy_db(BackupDir) of
         ok         -> info("upgrades: Mnesia dir backed up to ~p~n",
                            [BackupDir]);
@@ -127,7 +128,9 @@ remove_backup() ->
 maybe_upgrade_mnesia() ->
     AllNodes = rabbit_mnesia:cluster_nodes(all),
     ok = rabbit_mnesia_rename:maybe_finish(AllNodes),
-    case rabbit_version:upgrades_required(mnesia) of
+    %% Mnesia upgrade is the first upgrade scope,
+    %% so we should create a backup here if there are any upgrades
+    case rabbit_version:all_upgrades_required([mnesia, local, message_store]) of
         {error, starting_from_scratch} ->
             ok;
         {error, version_not_available} ->
@@ -143,10 +146,15 @@ maybe_upgrade_mnesia() ->
             ok;
         {ok, Upgrades} ->
             ensure_backup_taken(),
-            ok = case upgrade_mode(AllNodes) of
-                     primary   -> primary_upgrade(Upgrades, AllNodes);
-                     secondary -> secondary_upgrade(AllNodes)
-                 end
+            run_mnesia_upgrades(proplists:get_value(mnesia, Upgrades, []),
+                                AllNodes)
+    end.
+
+run_mnesia_upgrades([], _) -> ok;
+run_mnesia_upgrades(Upgrades, AllNodes) ->
+    case upgrade_mode(AllNodes) of
+        primary   -> primary_upgrade(Upgrades, AllNodes);
+        secondary -> secondary_upgrade(AllNodes)
     end.
 
 upgrade_mode(AllNodes) ->
@@ -244,24 +252,29 @@ maybe_upgrade_local() ->
         {ok, []}                       -> ensure_backup_removed(),
                                           ok;
         {ok, Upgrades}                 -> mnesia:stop(),
-                                          ensure_backup_taken(),
                                           ok = apply_upgrades(local, Upgrades,
                                                               fun () -> ok end),
-                                          ensure_backup_removed(),
                                           ok
     end.
 
 %% -------------------------------------------------------------------
 
 maybe_migrate_queues_to_per_vhost_storage() ->
-    case rabbit_version:upgrades_required(queues) of
+    Result = case rabbit_version:upgrades_required(message_store) of
         {error, version_not_available} -> version_not_available;
         {error, starting_from_scratch} -> starting_from_scratch;
         {error, _} = Err               -> throw(Err);
         {ok, []}                       -> ok;
-        {ok, Upgrades}                 -> apply_upgrades(queues, Upgrades,
-                                                         fun() -> ok end)
-    end.
+        {ok, Upgrades}                 -> apply_upgrades(message_store,
+                                                         Upgrades,
+                                                         fun() -> ok end),
+                                          ok
+    end,
+    %% Message store upgrades should be
+    %% the last group.
+    %% Backup can be deleted here.
+    ensure_backup_removed(),
+    Result.
 
 %% -------------------------------------------------------------------
 
diff --git a/src/rabbit_version.erl b/src/rabbit_version.erl
index a27f0aca00..4e2edd19eb 100644
--- a/src/rabbit_version.erl
+++ b/src/rabbit_version.erl
@@ -18,7 +18,8 @@
 
 -export([recorded/0, matches/2, desired/0, desired_for_scope/1,
          record_desired/0, record_desired_for_scope/1,
-         upgrades_required/1, check_version_consistency/3,
+         upgrades_required/1, all_upgrades_required/1,
+         check_version_consistency/3,
          check_version_consistency/4, check_otp_consistency/1,
          version_error/3]).
 
@@ -117,6 +118,30 @@ upgrades_required(Scope) ->
               end, Scope)
     end.
 
+all_upgrades_required(Scopes) ->
+    case recorded() of
+        {error, enoent} ->
+            case filelib:is_file(rabbit_guid:filename()) of
+                false -> {error, starting_from_scratch};
+                true  -> {error, version_not_available}
+            end;
+        {ok, _} ->
+            lists:foldl(
+                fun
+                (_, {error, Err}) -> {error, Err};
+                (Scope, {ok, Acc}) ->
+                    case upgrades_required(Scope) of
+                        %% Lift errors from any scope.
+                        {error, Err}   -> {error, Err};
+                        %% Filter non-upgradable scopes
+                        {ok, []}       -> {ok, Acc};
+                        {ok, Upgrades} -> {ok, [{Scope, Upgrades} | Acc]}
+                    end
+                end,
+                {ok, []},
+                Scopes)
+    end.
+
 %% -------------------------------------------------------------------
 
 with_upgrade_graph(Fun, Scope) ->