Repository: plataformatec/broadway
Pull request #59 (2019-03-26 19:13:39)

Allow producers to be canceled for shutdown
This is required for active producers like `BroadwayRabbitmq.Producer`

josevalim: Looks great. Just one comment about the function name and once we address it we should be good to go!
josevalim: @msaraiva looks good, just need docs!

Patch:
diff --git a/lib/broadway/producer.ex b/lib/broadway/producer.ex
index 8a0dc58..f23fd7a 100644
--- a/lib/broadway/producer.ex
+++ b/lib/broadway/producer.ex
@@ -3,6 +3,9 @@ defmodule Broadway.Producer do
   use GenStage
   alias Broadway.Message
 
+  @callback prepare_for_draining(state :: any) :: any
+  @optional_callbacks prepare_for_draining: 1
+
   @spec start_link(term, GenServer.options()) :: GenServer.on_start()
   def start_link(args, opts \\ []) do
     GenStage.start_link(__MODULE__, args, opts)
@@ -17,17 +20,24 @@ defmodule Broadway.Producer do
   def init(args) do
     {module, arg} = args[:module]
     transformer = args[:transformer]
-    # TODO: Raise a proper error message if we don't get  {:producer, state} back.
-    {:producer, module_state} = module.init(arg)
 
     state = %{
       module: module,
-      module_state: module_state,
+      module_state: nil,
       transformer: transformer,
       consumers: []
     }
 
-    {:producer, state}
+    case module.init(arg) do
+      {:producer, module_state} ->
+        {:producer, %{state | module_state: module_state}}
+
+      {:producer, module_state, options} ->
+        {:producer, %{state | module_state: module_state}, options}
+
+      return_value ->
+        {:stop, {:bad_return_value, return_value}}
+    end
   end
 
   @impl true
@@ -58,6 +68,17 @@ defmodule Broadway.Producer do
     end
   end
 
+  @impl true
+  def handle_cast(:prepare_for_draining, state) do
+    %{module: module, module_state: module_state} = state
+
+    if function_exported?(module, :prepare_for_draining, 1) do
+      module.prepare_for_draining(module_state)
+    end
+
+    {:noreply, [], state}
+  end
+
   @impl true
   def handle_info(:cancel_consumers, state) do
     for from <- state.consumers do
diff --git a/lib/broadway/terminator.ex b/lib/broadway/terminator.ex
index 80f7112..baeaaaa 100644
--- a/lib/broadway/terminator.ex
+++ b/lib/broadway/terminator.ex
@@ -41,6 +41,7 @@ defmodule Broadway.Terminator do
     end
 
     for name <- state.producers, pid = Process.whereis(name) do
+      GenStage.cast(pid, :prepare_for_draining)
       GenStage.demand(pid, :accumulate)
       GenStage.async_info(pid, :cancel_consumers)
     end


Test patch:
diff --git a/test/broadway/producer_test.exs b/test/broadway/producer_test.exs
index 5f01268..ef544c9 100644
--- a/test/broadway/producer_test.exs
+++ b/test/broadway/producer_test.exs
@@ -59,10 +59,21 @@ defmodule Broadway.ProducerTest do
     def init(_), do: {:producer, nil}
   end
 
+  defmodule ProducerWithBadReturn do
+    use GenStage
+
+    def init(_), do: {:consumer, nil}
+  end
+
   setup do
     %{state: %{module: FakeProducer, transformer: nil, module_state: nil}}
   end
 
+  test "init with bad return" do
+    args = %{module: {ProducerWithBadReturn, []}}
+    assert Producer.init(args) == {:stop, {:bad_return_value, {:consumer, nil}}}
+  end
+
   describe "wrap handle_demand" do
     test "returning {:noreply, [event], new_state}", %{state: state} do
       state = %{state | module_state: :return_no_reply}
diff --git a/test/broadway_test.exs b/test/broadway_test.exs
index a709c03..ad1e8b4 100644
--- a/test/broadway_test.exs
+++ b/test/broadway_test.exs
@@ -26,6 +26,8 @@ defmodule BroadwayTest do
   defmodule ManualProducer do
     use GenStage
 
+    @behaviour Broadway.Producer
+
     def start_link(args, opts \\ []) do
       GenStage.start_link(__MODULE__, args, opts)
     end
@@ -36,13 +38,40 @@ defmodule BroadwayTest do
       {:producer, %{test_pid: test_pid}}
     end
 
+    @impl true
     def init(_args) do
       {:producer, %{}}
     end
 
+    @impl true
     def handle_demand(_demand, state) do
       {:noreply, [], state}
     end
+
+    @impl true
+    def handle_info({:push_messages_async, messages}, state) do
+      {:noreply, messages, state}
+    end
+
+    @impl true
+    def prepare_for_draining(%{test_pid: test_pid}) do
+      message = wrap_message(:message_during_cancel, test_pid)
+      send(self(), {:push_messages_async, [message]})
+
+      message = wrap_message(:message_after_cancel, test_pid)
+      Process.send_after(self(), {:push_messages_async, [message]}, 1)
+      :ok
+    end
+
+    @impl true
+    def prepare_for_draining(_state) do
+      :ok
+    end
+
+    defp wrap_message(message, test_pid) do
+      ack = {CallerAcknowledger, {test_pid, make_ref()}, :ok}
+      %Message{data: message, acknowledger: ack}
+    end
   end
 
   defmodule EventProducer do
@@ -1039,7 +1068,7 @@ defmodule BroadwayTest do
         Broadway.start_link(CustomHandlers,
           name: broadway_name,
           producers: [
-            default: [module: {ManualProducer, []}]
+            default: [module: {ManualProducer, %{test_pid: self()}}]
           ],
           processors: [default: [stages: 1, min_demand: 1, max_demand: 4]],
           batchers: [default: [batch_size: 4]],
@@ -1084,6 +1113,15 @@ defmodule BroadwayTest do
       assert_receive {:ack, ^ref, [%{data: 1}, %{data: 2}], []}
     end
 
+    test "shutting down broadway cancels producers and waits for messages sent during cancellation",
+         %{broadway: broadway} do
+      Process.exit(broadway, :shutdown)
+
+      assert_receive {:ack, _, [%{data: :message_during_cancel}], []}
+      refute_receive {:ack, _, [%{data: :message_after_cancel}], []}
+      assert_receive {:EXIT, ^broadway, :shutdown}
+    end
+
     @tag shutdown: 1
     test "shutting down broadway respects shutdown value",
          %{broadway: broadway} do