Fix diameter_reg:subscribe/2 remove notification

Commit fae8ca0c broke notification by removing table elements before
matching for them, causing diameter_tcp/sctp listening processes to
live on after diameter:remove_transport/2.

Commit 58091992 added diameter_reg:subscribe/2, and commit 5ca5fb71
started using it for listener exit.

1 files changed, 11 insertions, 6 deletions

diff --git a/lib/diameter/src/base/diameter_reg.erl b/lib/diameter/src/base/diameter_reg.erl
index 5b7cfab31a..c1762a07e3 100644
--- a/lib/diameter/src/base/diameter_reg.erl
+++ b/lib/diameter/src/base/diameter_reg.erl
@@ -246,8 +246,11 @@ handle_call({add, Uniq, Key}, {Pid, _}, S) ->
 
 handle_call({remove, Key}, {Pid, _}, S) ->
     Rec = {Key, Pid},
-    ets:delete_object(?TABLE, Rec),
-    {reply, true, notify(remove, Rec, S)};
+    {reply, true, try
+                      notify(remove, Rec, S)
+                  after
+                      ets:delete_object(?TABLE, Rec)
+                  end};
 
 handle_call({wait, Pat}, {Pid, _} = From, S) ->
     NS = add_monitor(Pid, S),
@@ -370,10 +373,12 @@ send({_,_} = From, add, Rec) ->
 
 down(Pid, #state{monitors = Ps} = S) ->
     Recs = match('_', Pid),
-    ets:match_delete(?TABLE, {'_', Pid}),
-    lists:foldl(fun(R,NS) -> notify(remove, R, NS) end,
-                flush(Pid, S#state{monitors = sets:del_element(Pid, Ps)}),
-                Recs).
+    Acc0 = flush(Pid, S#state{monitors = sets:del_element(Pid, Ps)}),
+    try
+        lists:foldl(fun(R,NS) -> notify(remove, R, NS) end, Acc0, Recs)
+    after
+        ets:match_delete(?TABLE, {'_', Pid})
+    end.
 
 %% flush/3