2 files changed, 13 insertions, 10 deletions
diff --git a/lib/diameter/src/base/diameter_lib.erl b/lib/diameter/src/base/diameter_lib.erl
index 5b3a2063f8..b5f1ae3937 100644
--- a/lib/diameter/src/base/diameter_lib.erl
+++ b/lib/diameter/src/base/diameter_lib.erl
@@ -1,7 +1,7 @@
 %%
 %% %CopyrightBegin%
 %%
-%% Copyright Ericsson AB 2010-2014. All Rights Reserved.
+%% Copyright Ericsson AB 2010-2015. All Rights Reserved.
 %%
 %% The contents of this file are subject to the Erlang Public License,
 %% Version 1.1, (the "License"); you may not use this file except in
@@ -247,17 +247,19 @@ opts(HeapSize, Opts) ->
 %% # wait/1
 %% ---------------------------------------------------------------------------
 
--spec wait([pid()])
+-spec wait([pid() | reference()])
    -> ok.
 
 wait(L) ->
-    down([erlang:monitor(process, P) || P <- L]).
+    lists:foreach(fun down/1, L).
 
-down([]) ->
-    ok;
-down([MRef|T]) ->
-    receive {'DOWN', MRef, process, _, _} -> ok end,
-    down(T).
+down(Pid)
+  when is_pid(Pid) ->
+    down(monitor(process, Pid));
+
+down(MRef)
+  when is_reference(MRef) ->
+    receive {'DOWN', MRef, process, _, _} = T -> T end.
 
 %% ---------------------------------------------------------------------------
 %% # fold_tuple/3
diff --git a/lib/diameter/src/base/diameter_service.erl b/lib/diameter/src/base/diameter_service.erl
index 76b05a2ad4..f4afa6eb94 100644
--- a/lib/diameter/src/base/diameter_service.erl
+++ b/lib/diameter/src/base/diameter_service.erl
@@ -1,7 +1,7 @@
 %%
 %% %CopyrightBegin%
 %%
-%% Copyright Ericsson AB 2010-2014. All Rights Reserved.
+%% Copyright Ericsson AB 2010-2015. All Rights Reserved.
 %%
 %% The contents of this file are subject to the Erlang Public License,
 %% Version 1.1, (the "License"); you may not use this file except in
@@ -610,8 +610,9 @@ st(#watchdog{ref = Ref, pid = Pid}, Refs) ->
 %% st/3
 
 st(#watchdog{pid = Pid}, Reason, Acc) ->
+    MRef = monitor(process, Pid),
     Pid ! {shutdown, self(), Reason},
-    [Pid | Acc].
+    [MRef | Acc].
 
 %% ---------------------------------------------------------------------------
 %% # call_service/2