Do not require {ref, :done}, instead allow sync removal

Signed-off-by: José Valim <[email protected]>

Author: José Valim

lib/elixir/lib/gen_event.ex

@@ -188,7 +188,7 @@ defmodule GenEvent do
   @type manager :: pid | name | {atom, node}
 
   @typedoc "Supported values for new handlers"
-  @type handler :: atom | {atom, term} | pid
+  @type handler :: atom | {atom, term} | {pid, reference}
 
   @doc false
   defmacro __using__(_) do
@@ -612,8 +612,8 @@ defmodule GenEvent do
         {hib, reply, handlers} = server_add_mon_handler(handler, args, handlers, notify)
         reply(tag, reply)
         loop(parent, name, handlers, debug, hib)
-      {_from, tag, {:add_process_handler, pid, notify}} ->
-        {hib, reply, handlers} = server_add_process_handler(pid, handlers, notify)
+      {_from, tag, {:add_process_handler, pid, notify, ref}} ->
+        {hib, reply, handlers} = server_add_process_handler(pid, handlers, notify, ref)
         reply(tag, reply)
         loop(parent, name, handlers, debug, hib)
       {_from, tag, {:delete_handler, handler, args}} ->
@@ -761,9 +761,9 @@ defmodule GenEvent do
     do_add_handler(module, handler, args, handlers, :ok)
   end
 
-  defp server_add_process_handler(pid, handlers, notify) do
-    ref = Process.monitor(pid)
-    handler = handler(module: GenEvent.Stream, id: pid,
+  defp server_add_process_handler(pid, handlers, notify, ref) do
+    ref = ref || Process.monitor(pid)
+    handler = handler(module: GenEvent.Stream, id: {self(), ref},
                       pid: notify, ref: ref)
     do_add_handler(GenEvent.Stream, handler, {pid, ref}, handlers, {self(), ref})
   end
@@ -796,7 +796,7 @@ defmodule GenEvent do
 
   defp server_split_process_handlers(mode, event, [handler|t], handlers, streams) do
     case handler(handler, :id) do
-      pid when is_pid(pid) ->
+      {pid, _ref} when is_pid(pid) ->
         server_process_notify(mode, event, handler)
         server_split_process_handlers(mode, event, t, handlers, [handler|streams])
       _ ->
@@ -858,16 +858,17 @@ defmodule GenEvent do
   end
 
   defp server_collect_process_handlers(mode, event, [handler|t], handlers, name) when mode in [:sync, :ack] do
-    handler(ref: ref) = handler
+    handler(ref: ref, id: id) = handler
 
     receive do
       {^ref, :ok} ->
         server_collect_process_handlers(mode, event, t, [handler|handlers], name)
-      {^ref, :done} ->
-        do_terminate(handler, :remove_handler, event, name, :normal)
+      {_from, tag, {:delete_handler, ^id, args}} ->
+        do_terminate(handler, args, :remove, name, :normal)
+        reply(tag, :ok)
         server_collect_process_handlers(mode, event, t, handlers, name)
       {:DOWN, ^ref, _, _, reason} ->
-        do_terminate(handler, {:stop, reason}, :remove, name, :shutdown)
+        do_terminate(handler, {:stop, reason}, :DOWN, name, :shutdown)
         server_collect_process_handlers(mode, event, t, handlers, name)
     end
   end
@@ -939,7 +940,7 @@ defmodule GenEvent do
     case :lists.keyfind(ref, handler(:ref) + 1, handlers) do
       false -> :error
       handler ->
-        do_terminate(handler, {:stop, reason}, :remove, name, :shutdown)
+        do_terminate(handler, {:stop, reason}, :DOWN, name, :shutdown)
         {:ok, :lists.keydelete(ref, handler(:ref) + 1, handlers)}
     end
   end

lib/elixir/lib/gen_event/stream.ex

@@ -80,7 +80,8 @@ defimpl Enumerable, for: GenEvent.Stream do
 
   defp start(%{manager: manager} = stream) do
     try do
-      {:ok, {pid, ref}} = :gen.call(manager, self(), {:add_process_handler, self(), self()}, :infinity)
+      {:ok, {pid, ref}} = :gen.call(manager, self(),
+                                    {:add_process_handler, self(), self(), nil}, :infinity)
       mon_ref = Process.monitor(pid)
       {pid, ref, mon_ref}
     catch
@@ -92,19 +93,6 @@ defimpl Enumerable, for: GenEvent.Stream do
     self = self()
 
     receive do
-      # The handler was removed. Stop iteration, resolve the
-      # event later. We need to demonitor now, otherwise DOWN
-      # appears with higher priority in the shutdown process.
-      {:gen_event_EXIT, ^self, _reason} = event ->
-        Process.demonitor(mon_ref, [:flush])
-        send(self, event)
-        {:halt, {:removed, acc}}
-
-      # The manager died. Stop iteration, resolve the event later.
-      {:DOWN, ^mon_ref, _, _, _} = event ->
-        send(self, event)
-        {:halt, {:removed, acc}}
-
       # Got an async event.
       {_from, {^pid, ^ref}, {:notify, event}} ->
         {[{:async, pid, ref, event}], acc}
@@ -116,6 +104,19 @@ defimpl Enumerable, for: GenEvent.Stream do
       # Got an ack event.
       {_from, {^pid, ^ref}, {:ack_notify, event}} ->
         {[{:ack, pid, ref, event}], acc}
+
+      # The handler was removed. Stop iteration, resolve the
+      # event later. We need to demonitor now, otherwise DOWN
+      # appears with higher priority in the shutdown process.
+      {:gen_event_EXIT, {^pid, ^ref}, _reason} = event ->
+        Process.demonitor(mon_ref, [:flush])
+        send(self, event)
+        {:halt, {:removed, acc}}
+
+      # The manager died. Stop iteration, resolve the event later.
+      {:DOWN, ^mon_ref, _, _, _} = event ->
+        send(self, event)
+        {:halt, {:removed, acc}}
     after
       timeout ->
         exit({:timeout, {__MODULE__, :next, [stream, acc]}})
@@ -135,26 +136,20 @@ defimpl Enumerable, for: GenEvent.Stream do
 
   # If we reach this branch, the handler was not removed yet,
   # so we trigger a request for doing so.
-  defp stop(stream, {pid, _, _} = acc) do
-    parent = self()
-    _ = Task.start(fn -> GenEvent.remove_handler(pid, parent, :shutdown) end)
+  defp stop(stream, {pid, ref, _} = acc) do
+    _ = GenEvent.remove_handler(pid, {pid, ref}, :shutdown)
     stop(stream, {:removed, acc})
   end
 
   defp wait_for_handler_removal(pid, ref, mon_ref) do
-    self = self()
-
     receive do
-      {_from, {^pid, ^ref}, {notify, _event}} when notify in [:ack_notify, :sync_notify] ->
-        send pid, {ref, :done}
-        wait_for_handler_removal(pid, ref, mon_ref)
-      {:gen_event_EXIT, ^self, reason}
+      {:gen_event_EXIT, {^pid, ^ref}, reason}
           when reason == :normal
           when reason == :shutdown
           when tuple_size(reason) == 3 and elem(reason, 0) == :swapped  ->
         Process.demonitor(mon_ref, [:flush])
         :ok
-      {:gen_event_EXIT, ^self, reason} ->
+      {:gen_event_EXIT, {^pid, ^ref}, reason} ->
         Process.demonitor(mon_ref, [:flush])
         {:error, reason}
       {:DOWN, ^mon_ref, _, _, reason} ->
@@ -164,7 +159,7 @@ defimpl Enumerable, for: GenEvent.Stream do
 
   defp flush_events(ref) do
     receive do
-      {_from, {_pid, ^ref}, {:notify, _event}} ->
+      {_from, {_pid, ^ref}, {notify, _event}} when notify in [:notify, :ack_notify, :sync_notify] ->
         flush_events(ref)
     after
       0 -> :ok