Avoid sending events in accumulate demand (#296)

Author: manuel-rubio

lib/gen_stage.ex

@@ -1339,9 +1339,9 @@ defmodule GenStage do
   Sets the demand mode for a producer.
 
   When `:forward`, the demand is always forwarded to the `c:handle_demand/2`
-  callback. When `:accumulate`, demand is accumulated until its mode is
-  set to `:forward`. This is useful as a synchronization mechanism, where
-  the demand is accumulated until all consumers are subscribed. Defaults
+  callback. When `:accumulate`, both demand and events are accumulated until
+  its mode is set to `:forward`. This is useful as a synchronization mechanism,
+  where the demand is accumulated until all consumers are subscribed. Defaults
   to `:forward`.
 
   This command is asynchronous.
@@ -2261,11 +2261,11 @@ defmodule GenStage do
 
     if is_list(events) do
       fold_fun = fn
-        d, {:noreply, %{state: state} = stage} ->
-          noreply_callback(:handle_demand, [d, state], stage)
+        event, {:noreply, stage} ->
+          handle_accumulated_event(event, stage)
 
-        d, {:noreply, %{state: state} = stage, _} ->
-          noreply_callback(:handle_demand, [d, state], stage)
+        event, {:noreply, stage, _} ->
+          handle_accumulated_event(event, stage)
 
         _, {:stop, _, _} = acc ->
           acc
@@ -2285,6 +2285,14 @@ defmodule GenStage do
     end
   end
 
+  defp handle_accumulated_event({:demand, d}, stage) do
+    take_from_buffer_or_handle_demand(d, stage)
+  end
+
+  defp handle_accumulated_event({:dispatch, events, length}, stage) do
+    {:noreply, dispatch_events(events, length, stage)}
+  end
+
   defp producer_subscribe(opts, from, stage) do
     %{mod: mod, state: state, dispatcher_mod: dispatcher_mod, dispatcher_state: dispatcher_state} =
       stage
@@ -2370,23 +2378,33 @@ defmodule GenStage do
         take_pc_events(queue, counter, stage)
 
       %{} ->
-        case take_from_buffer(counter, %{stage | dispatcher_state: dispatcher_state}) do
-          {:ok, 0, stage} ->
-            {:noreply, stage}
+        take_from_buffer_or_handle_demand(counter, %{stage | dispatcher_state: dispatcher_state})
+    end
+  end
 
-          {:ok, counter, %{events: :forward, state: state} = stage} ->
-            noreply_callback(:handle_demand, [counter, state], stage)
+  defp take_from_buffer_or_handle_demand(counter, stage) do
+    case take_from_buffer(counter, stage) do
+      {:ok, 0, stage} ->
+        {:noreply, stage}
 
-          {:ok, counter, %{events: events} = stage} when is_list(events) ->
-            {:noreply, %{stage | events: [counter | events]}}
-        end
+      {:ok, counter, %{events: :forward, state: state} = stage} ->
+        noreply_callback(:handle_demand, [counter, state], stage)
+
+      {:ok, counter, %{events: events} = stage} when is_list(events) ->
+        {:noreply, %{stage | events: [{:demand, counter} | events]}}
     end
   end
 
   defp dispatch_events([], _length, stage) do
     stage
   end
 
+  # We don't dispatch when we are accumulating demand
+  defp dispatch_events(to_dispatch, length, %{events: events, type: :producer} = stage)
+       when is_list(events) do
+    %{stage | events: [{:dispatch, to_dispatch, length} | events]}
+  end
+
   defp dispatch_events(events, _length, %{type: :consumer} = stage) do
     error_msg =
       ~c"GenStage consumer ~tp cannot dispatch events (an empty list must be returned): ~tp~n"

test/gen_stage_test.exs

@@ -3,6 +3,29 @@ defmodule GenStageTest do
 
   import ExUnit.CaptureLog
 
+  defmodule EventProducer do
+    @moduledoc """
+    Produce events when receives a cast
+    """
+    use GenStage
+
+    def start_link(init) do
+      GenStage.start_link(__MODULE__, init)
+    end
+
+    @impl GenStage
+    @doc false
+    def init(init), do: init
+
+    @impl GenStage
+    @doc false
+    def handle_cast(event, state), do: {:noreply, [event], state}
+
+    @impl GenStage
+    @doc false
+    def handle_demand(_demand, state), do: {:noreply, [], state}
+  end
+
   defmodule Counter do
     @moduledoc """
     A producer that works as a counter in batches.
@@ -673,6 +696,38 @@ defmodule GenStageTest do
       assert_receive {:consumed, [0, 1, 2, 3]}
     end
 
+    test "can be set to :accumulate via API using broadcast" do
+      {:ok, producer} =
+        EventProducer.start_link({:producer, [], dispatcher: GenStage.BroadcastDispatcher})
+
+      assert GenStage.demand(producer) == :forward
+      {:ok, consumer1} = Forwarder.start_link({:consumer, self(), subscribe_to: [producer]})
+      {:ok, consumer2} = Forwarder.start_link({:consumer, self(), subscribe_to: [producer]})
+      GenStage.demand(producer, :accumulate)
+      assert GenStage.demand(producer) == :accumulate
+
+      GenStage.stop(consumer1)
+      GenStage.stop(consumer2)
+
+      assert :ok = GenStage.cast(producer, 1)
+      assert :ok = GenStage.cast(producer, 2)
+      assert :ok = GenStage.cast(producer, 3)
+      assert :ok = GenStage.cast(producer, 4)
+
+      Process.sleep(200)
+      {:ok, _consumer1} = Forwarder.start_link({:consumer, self(), subscribe_to: [producer]})
+      {:ok, _consumer2} = Forwarder.start_link({:consumer, self(), subscribe_to: [producer]})
+      refute_receive {:consumed, _}
+
+      GenStage.demand(producer, :forward)
+      assert GenStage.demand(producer) == :forward
+      assert {{[], []}, 0, _} = :sys.get_state(producer).buffer
+      assert_receive {:consumed, [1]}
+      assert_receive {:consumed, [2]}
+      assert_receive {:consumed, [3]}
+      assert_receive {:consumed, [4]}
+    end
+
     test "can be set to :accumulate via API" do
       {:ok, producer} = Counter.start_link({:producer, 0})
       assert GenStage.demand(producer) == :forward