Aggregate which_children, implement count_children.

2 files changed, 25 insertions, 2 deletions

diff --git a/src/mirrored_supervisor.erl b/src/mirrored_supervisor.erl
index eb987189f8..7f99ac6fd8 100644
--- a/src/mirrored_supervisor.erl
+++ b/src/mirrored_supervisor.erl
@@ -118,7 +118,7 @@
 -export([start_link/3, start_link/4,
          start_child/2, restart_child/2,
          delete_child/2, terminate_child/2,
-         which_children/1, check_childspecs/1]).
+         which_children/1, count_children/1, check_childspecs/1]).
 
 -export([behaviour_info/1]).
 
@@ -259,7 +259,8 @@ start_child(Sup, ChildSpec) -> call(Sup, {start_child,  ChildSpec}).
 delete_child(Sup, Id)       -> find_call(Sup, Id, {delete_child, Id}).
 restart_child(Sup, Id)      -> find_call(Sup, Id, {msg, restart_child, [Id]}).
 terminate_child(Sup, Id)    -> find_call(Sup, Id, {msg, terminate_child, [Id]}).
-which_children(Sup)         -> ?SUPERVISOR:which_children(child(Sup, delegate)).
+which_children(Sup)         -> fold(which_children, Sup, fun lists:append/2).
+count_children(Sup)         -> fold(count_children, Sup, fun add_proplists/2).
 check_childspecs(Specs)     -> ?SUPERVISOR:check_childspecs(Specs).
 
 behaviour_info(callbacks) -> [{init,1}];
@@ -281,6 +282,13 @@ find_call(Sup, Id, Msg) ->
         []       -> {error, not_found}
     end.
 
+fold(FunAtom, Sup, AggFun) ->
+    Group = call(Sup, group),
+    lists:foldl(AggFun, [],
+                [apply(?SUPERVISOR, FunAtom, [D]) ||
+                    M <- ?PG2:get_members(Group),
+                    D <- [?GEN_SERVER:call(M, delegate_supervisor, infinity)]]).
+
 child(Sup, Id) ->
     [Pid] = [Pid || {Id1, Pid, _, _} <- ?SUPERVISOR:which_children(Sup),
                     Id1 =:= Id],
@@ -517,3 +525,10 @@ with_exit_handler(Handler, Thunk) ->
         exit:{{R, _}, _} when R =:= nodedown; R =:= shutdown ->
             Handler()
     end.
+
+add_proplists(P1, []) -> P1;
+add_proplists(P1, [{K, V2} | P2]) ->
+    add_proplists(case proplists:get_value(K, P1) of
+                      undefined -> [{K, V2} | P1];
+                      V1        -> [{K, V1 + V2} | proplists:delete(K, P1)]
+                  end, P2).
diff --git a/src/mirrored_supervisor_tests.erl b/src/mirrored_supervisor_tests.erl
index 5a48b2221c..aea9c05c53 100644
--- a/src/mirrored_supervisor_tests.erl
+++ b/src/mirrored_supervisor_tests.erl
@@ -35,6 +35,7 @@ all_tests() ->
     passed = test_migrate_twice(),
     passed = test_already_there(),
     passed = test_delete_restart(),
+    passed = test_which_children(),
     passed = test_large_group(),
     passed = test_childspecs_at_init(),
     passed = test_anonymous_supervisors(),
@@ -97,6 +98,13 @@ test_delete_restart() ->
                       false = (Pid3 =:= Pid4)
               end, [a, b]).
 
+test_which_children() ->
+    with_sups(fun([A, B]) ->
+                      ?MS:start_child(A, childspec(worker)),
+                      ?MS:start_child(B, childspec(worker2)),
+                      2 = length(?MS:which_children(A))
+              end, [a, b]).
+
 %% Not all the members of the group should actually do the failover
 test_large_group() ->
     with_sups(fun([A, _, _, _]) ->