After rebase fixes

Author: Daniil Fedotov

src/rabbit_variable_queue.erl

@@ -620,7 +620,8 @@ terminate(_Reason, State) ->
         purge_pending_ack(true, State),
     PRef = case MSCStateP of
                undefined -> undefined;
-               {MP, CP}  -> ok = MP:client_terminate(CP),
+               {MP, CP}  ->
+                            ok = maybe_client_terminate(MP, CP),
                             MP:client_ref(CP)
            end,
     {MT, CT} = MSCStateT,
@@ -2838,7 +2839,7 @@ move_messages_to_vhost_store(Queues) ->
     ok = delete_old_store(OldStore),
     ok = rabbit_queue_index:cleanup_global_recovery_terms(),
     [ok= rabbit_recovery_terms:close_table(VHost) || VHost <- VHosts],
-    ok = stop_new_store(NewMsgStore).
+    ok = stop_new_store(NewMsgStore),
     rabbit_file:write_term_file(msg_store_module_file(), [rabbit_msg_store]),
     ok.
 
@@ -2992,12 +2993,12 @@ log_upgrade_verbose(Msg) ->
 log_upgrade_verbose(Msg, Args) ->
     rabbit_log_upgrade:info(Msg, Args).
 
-maybe_client_terminate(MSCStateP) ->
+maybe_client_terminate(MP, CP) ->
     %% Queue might have been asked to stop by the supervisor, it needs a clean
     %% shutdown in order for the supervising strategy to work - if it reaches max
     %% restarts might bring the vhost down.
     try
-        rabbit_msg_store:client_terminate(MSCStateP)
+        MP:client_terminate(CP)
     catch
         _:_ ->
             ok

src/rabbit_vhost_msg_store.erl

@@ -21,14 +21,16 @@
 -export([start/4, stop/2, client_init/5, successfully_recovered_state/2]).
 
 
-start(VHost, Type, ClientRefs, StartupFunState) when is_list(ClientRefs);
-                                                     ClientRefs == undefined  ->
+start(VHost, Type, ClientRefs, StartupFunState, MsgStoreModule)
+when is_list(ClientRefs); ClientRefs == undefined  ->
     case rabbit_vhost_sup_sup:vhost_sup(VHost) of
         {ok, VHostSup} ->
+            ets:insert(rabbit_vhost_sup_sup, {{VHost, Type}, MsgStoreModule}),
             VHostDir = rabbit_vhost:msg_store_dir_path(VHost),
             supervisor2:start_child(VHostSup,
                                     {Type, {rabbit_msg_store, start_link,
-                                            [Type, VHostDir, ClientRefs, StartupFunState]},
+                                            [Type, VHostDir, ClientRefs,
+                                             StartupFunState, MsgStoreModule]},
                                      transient, ?WORKER_WAIT, worker, [rabbit_msg_store]});
         %% we can get here if a vhost is added and removed concurrently
         %% e.g. some integration tests do it
@@ -52,16 +54,19 @@ stop(VHost, Type) ->
     end.
 
 client_init(VHost, Type, Ref, MsgOnDiskFun, CloseFDsFun) ->
-    with_vhost_store(VHost, Type, fun(StorePid) ->
-        rabbit_msg_store:client_init(StorePid, Ref, MsgOnDiskFun, CloseFDsFun)
+    with_vhost_store(VHost, Type, fun(StorePid, MsgStoreModule) ->
+        MsgStoreModule:client_init(StorePid, Ref, MsgOnDiskFun, CloseFDsFun)
     end).
 
 with_vhost_store(VHost, Type, Fun) ->
     case vhost_store_pid(VHost, Type) of
         no_pid ->
             throw({message_store_not_started, Type, VHost});
         Pid when is_pid(Pid) ->
-            Fun(Pid)
+            case ets:lookup(rabbit_vhost_sup_sup, {VHost, Type}) of
+                [MsgStoreModule] -> Fun(Pid, MsgStoreModule);
+                [] -> error({message_store_module_not_found, {VHost, Type, Pid}})
+            end
     end.
 
 vhost_store_pid(VHost, Type) ->
@@ -72,6 +77,6 @@ vhost_store_pid(VHost, Type) ->
     end.
 
 successfully_recovered_state(VHost, Type) ->
-    with_vhost_store(VHost, Type, fun(StorePid) ->
-        rabbit_msg_store:successfully_recovered_state(StorePid)
+    with_vhost_store(VHost, Type, fun(StorePid, MsgStoreModule) ->
+        MsgStoreModule:successfully_recovered_state(StorePid)
     end).