Queue and msg store supervisor locations

1 files changed, 51 insertions, 28 deletions

diff --git a/src/rabbit_vm.erl b/src/rabbit_vm.erl
index 24398c477b..1039ed659c 100644
--- a/src/rabbit_vm.erl
+++ b/src/rabbit_vm.erl
@@ -41,17 +41,18 @@ memory() ->
         [aggregate(Names, Sums, memory, fun (X) -> X end)
          || Names <- distinguished_interesting_sups()],
 
-    Mnesia              = mnesia_memory(),
-    MsgIndexETS  = ets_memory([msg_store_persistent_vhost, msg_store_transient_vhost]),
-    MetricsETS          = ets_memory([rabbit_metrics]),
-    MetricsProc  = try
-		       [{_, M}] = process_info(whereis(rabbit_metrics), [memory]),
-		       M
-		   catch
-		       error:badarg ->
-			   0
-		   end,
-    MgmtDbETS           = ets_memory([rabbit_mgmt_storage]),
+    Mnesia       = mnesia_memory(),
+    MsgIndexETS  = ets_memory(msg_stores()),
+    MetricsETS   = ets_memory([rabbit_metrics]),
+    MetricsProc  =
+        try
+            [{_, M}] = process_info(whereis(rabbit_metrics), [memory]),
+            M
+        catch
+            error:badarg ->
+            0
+        end,
+    MgmtDbETS    = ets_memory([rabbit_mgmt_storage]),
 
     [{total,     Total},
      {processes, Processes},
@@ -124,8 +125,8 @@ mnesia_memory() ->
         _   -> 0
     end.
 
-ets_memory(OwnerNames) ->
-    lists:sum([V || {_K, V} <- ets_tables_memory(OwnerNames)]).
+ets_memory(Owners) ->
+    lists:sum([V || {_K, V} <- ets_tables_memory(Owners)]).
 
 ets_tables_memory(all) ->
     [{ets:info(T, name), bytes(ets:info(T, memory))}
@@ -133,23 +134,44 @@ ets_tables_memory(all) ->
         is_atom(T)];
 ets_tables_memory(OwnerName) when is_atom(OwnerName) ->
     ets_tables_memory([OwnerName]);
-ets_tables_memory(OwnerNames) when is_list(OwnerNames) ->
-    Owners = [whereis(N) || N <- OwnerNames],
+ets_tables_memory(Owners) when is_list(Owners) ->
+    OwnerPids = lists:map(fun(O) when is_pid(O) -> O;
+                             (O) when is_atom(O) -> whereis(O)
+                          end,
+                          Owners),
     [{ets:info(T, name), bytes(ets:info(T, memory))}
      || T <- ets:all(),
-        lists:member(ets:info(T, owner), Owners)].
+        lists:member(ets:info(T, owner), OwnerPids)].
 
 bytes(Words) ->  try
                      Words * erlang:system_info(wordsize)
                  catch
                      _:_ -> 0
                  end.
-%% TODO: per-vhost supervisor
+
 interesting_sups() ->
-    [[rabbit_amqqueue_sup_sup], conn_sups() | interesting_sups0()].
+    [queue_sups(), conn_sups() | interesting_sups0()].
+
+queue_sups() ->
+    all_vhosts_children(rabbit_amqqueue_sup_sup).
+
+msg_stores() ->
+    all_vhosts_children(msg_store_transient)
+    ++
+    all_vhosts_children(msg_store_persistent).
+
+all_vhosts_children(Name) ->
+    lists:filter_map(
+        fun({_, VHostSup, _, _}) ->
+            case supervisor2:find_child(VHostSup, Name) of
+                [QSup] -> {true, QSup};
+                []     -> false
+            end
+        end,
+        supervisor:which_children(rabbit_vhost_sup_sup)).
 
 interesting_sups0() ->
-    MsgIndexProcs = [msg_store_transient_vhost, msg_store_persistent_vhost],
+    MsgIndexProcs = msg_stores(),
     MgmtDbProcs   = [rabbit_mgmt_sup_sup],
     PluginProcs   = plugin_sups(),
     [MsgIndexProcs, MgmtDbProcs, PluginProcs].
@@ -166,18 +188,19 @@ ranch_server_sups() ->
         error:badarg  -> []
     end.
 
-conn_sups(With) -> [{Sup, With} || Sup <- conn_sups()].
+with(Sups, With) -> [{Sup, With} || Sup <- Sups].
 
-distinguishers() -> [{rabbit_amqqueue_sup_sup, fun queue_type/1} |
-                     conn_sups(fun conn_type/1)].
+distinguishers() -> with(queue_sups(), fun queue_type/1) ++
+                    with(conn_sups(), fun conn_type/1).
 
 distinguished_interesting_sups() ->
-    [[{rabbit_amqqueue_sup_sup, master}],
-     [{rabbit_amqqueue_sup_sup, slave}],
-     conn_sups(reader),
-     conn_sups(writer),
-     conn_sups(channel),
-     conn_sups(other)]
+    [
+     with(queue_sups(), master),
+     with(queue_sups(), slave),
+     with(conn_sups(), reader),
+     with(conn_sups(), writer),
+     with(conn_sups(), channel),
+     with(conn_sups(), other)]
         ++ interesting_sups0().
 
 plugin_sups() ->