Make broadcast dispatcher subscribe idempotent (#225)

Author: evnu

lib/gen_stage.ex

@@ -1901,7 +1901,14 @@ defmodule GenStage do
             mon_ref = Process.monitor(consumer_pid)
             stage = put_in(stage.monitors[mon_ref], ref)
             stage = put_in(stage.consumers[ref], {consumer_pid, mon_ref})
-            producer_subscribe(opts, from, stage)
+
+            case producer_subscribe(opts, from, stage) do
+              {:error, :already_subscribed} ->
+                producer_cancel(ref, :cancel, :already_subscribed, stage)
+
+              other ->
+                other
+            end
 
           other ->
             other
@@ -2221,8 +2228,14 @@ defmodule GenStage do
   end
 
   defp dispatcher_callback(callback, args, %{dispatcher_mod: dispatcher_mod} = stage) do
-    {:ok, counter, dispatcher_state} = apply(dispatcher_mod, callback, args)
+    dispatcher_mod |> apply(callback, args) |> handle_callback_result(stage)
+  end
+
+  defp handle_callback_result(e = {:error, _reason}, _stage) do
+    e
+  end
 
+  defp handle_callback_result({:ok, counter, dispatcher_state}, stage) do
     case stage do
       %{type: :producer_consumer, events: {queue, demand}} ->
         counter = demand + counter
@@ -2500,11 +2513,17 @@ defmodule GenStage do
       {:noreply, stage}
       when mode == :permanent
       when mode == :transient and not Utils.is_transient_shutdown(reason) ->
-        error_msg =
-          'GenStage consumer ~tp is stopping after receiving cancel from producer ~tp with reason: ~tp~n'
+        case reason do
+          :already_subscribed ->
+            {:noreply, stage}
 
-        :error_logger.info_msg(error_msg, [Utils.self_name(), pid, reason])
-        {:stop, reason, stage}
+          _other ->
+            error_msg =
+              'GenStage consumer ~tp is stopping after receiving cancel from producer ~tp with reason: ~tp~n'
+
+            :error_logger.info_msg(error_msg, [Utils.self_name(), pid, reason])
+            {:stop, reason, stage}
+        end
 
       other ->
         other

lib/gen_stage/dispatcher.ex

@@ -41,7 +41,7 @@ defmodule GenStage.Dispatcher do
   Called every time the producer gets a new subscriber.
   """
   @callback subscribe(opts :: keyword(), from :: {pid, reference}, state :: term) ::
-              {:ok, demand :: non_neg_integer, new_state}
+              {:ok, demand :: non_neg_integer, new_state} | {:error, term}
             when new_state: term
 
   @doc """

lib/gen_stage/dispatchers/broadcast_dispatcher.ex

@@ -33,9 +33,11 @@ defmodule GenStage.BroadcastDispatcher do
 
   @behaviour GenStage.Dispatcher
 
+  require Logger
+
   @doc false
   def init(_opts) do
-    {:ok, {[], 0}}
+    {:ok, {[], 0, MapSet.new()}}
   end
 
   @doc false
@@ -45,36 +47,48 @@ defmodule GenStage.BroadcastDispatcher do
   end
 
   @doc false
-  def subscribe(opts, {pid, ref}, {demands, waiting}) do
+  def subscribe(opts, {pid, ref}, {demands, waiting, subscribed_processes}) do
     selector = validate_selector(opts)
-    {:ok, 0, {add_demand(-waiting, pid, ref, selector, demands), waiting}}
+
+    if subscribed?(subscribed_processes, pid) do
+      Logger.error(fn ->
+        "#{inspect(pid)} is already registered with #{inspect(self())}. " <>
+          "This subscription has been discared."
+      end)
+
+      {:error, :already_subscribed}
+    else
+      subscribed_processes = add_subscriber(subscribed_processes, pid)
+      {:ok, 0, {add_demand(-waiting, pid, ref, selector, demands), waiting, subscribed_processes}}
+    end
   end
 
   @doc false
-  def cancel({_, ref}, {demands, waiting}) do
+  def cancel({pid, ref}, {demands, waiting, subscribed_processes}) do
     # Since we may have removed the process we were waiting on,
     # cancellation may actually generate demand!
     demands = delete_demand(ref, demands)
     new_min = get_min(demands)
     demands = adjust_demand(new_min, demands)
-    {:ok, new_min, {demands, waiting + new_min}}
+    subscribed_processes = delete_subscriber(subscribed_processes, pid)
+    {:ok, new_min, {demands, waiting + new_min, subscribed_processes}}
   end
 
   @doc false
-  def ask(counter, {pid, ref}, {demands, waiting}) do
+  def ask(counter, {pid, ref}, {demands, waiting, subscribed_processes}) do
     {current, selector, demands} = pop_demand(ref, demands)
     demands = add_demand(current + counter, pid, ref, selector, demands)
     new_min = get_min(demands)
     demands = adjust_demand(new_min, demands)
-    {:ok, new_min, {demands, waiting + new_min}}
+    {:ok, new_min, {demands, waiting + new_min, subscribed_processes}}
   end
 
   @doc false
-  def dispatch(events, _length, {demands, 0}) do
-    {:ok, events, {demands, 0}}
+  def dispatch(events, _length, {demands, 0, subscribed_processes}) do
+    {:ok, events, {demands, 0, subscribed_processes}}
   end
 
-  def dispatch(events, length, {demands, waiting}) do
+  def dispatch(events, length, {demands, waiting, subscribed_processes}) do
     {deliver_now, deliver_later, waiting} = split_events(events, length, waiting)
 
     for {_, pid, ref, selector} <- demands do
@@ -92,7 +106,7 @@ defmodule GenStage.BroadcastDispatcher do
       :ok
     end
 
-    {:ok, deliver_later, {demands, waiting}}
+    {:ok, deliver_later, {demands, waiting, subscribed_processes}}
   end
 
   defp filter_and_count(messages, nil) do
@@ -168,4 +182,16 @@ defmodule GenStage.BroadcastDispatcher do
   defp delete_demand(ref, demands) do
     List.keydelete(demands, ref, 2)
   end
+
+  defp add_subscriber(subscribed_processes, pid) do
+    MapSet.put(subscribed_processes, pid)
+  end
+
+  defp delete_subscriber(subscribed_processes, pid) do
+    MapSet.delete(subscribed_processes, pid)
+  end
+
+  defp subscribed?(subscribed_processes, pid) do
+    MapSet.member?(subscribed_processes, pid)
+  end
 end