Repository: plataformatec/broadway
Pull request #113 (2019-09-22 08:22:18)

Inject topology index when starting producers
This PR injects a `:broadway_topology_index` option in the producer-specific options, but only if those options are a keyword list. I picked the first name I could think of, but not sure it's a great name so very open to renaming it :)


Patch:
diff --git a/lib/broadway.ex b/lib/broadway.ex
index 5cbfdd5..c5e3e29 100644
--- a/lib/broadway.ex
+++ b/lib/broadway.ex
@@ -460,6 +460,11 @@ defmodule Broadway do
       that implements the GenStage behaviour and `arg` the argument that will
       be passed to the `init/1` callback of the producer. Pay attention that
       this producer must emit events that are `Broadway.Message` structs.
+      It's recommended that `arg` is a keyword list. In fact, if `arg` is
+      a keyword list, a `:broadway_index` option is injected into
+      such keyword list: this index can be used to customize the behaviour
+      of a producer based on its index (for example, having even producers
+      connect to some server while odd producers connect to another).
     * `:stages` - Optional. The number of stages that will be
       created by Broadway. Use this option to control the concurrency
       level of each set of producers. The default value is `1`.
diff --git a/lib/broadway/producer.ex b/lib/broadway/producer.ex
index 6dd5954..791c6ca 100644
--- a/lib/broadway/producer.ex
+++ b/lib/broadway/producer.ex
@@ -11,7 +11,7 @@ defmodule Broadway.Producer do
   started before the ProducerSupervisor in Broadwday's supervision tree.
   """
   @callback prepare_for_start(module :: atom, options :: keyword) ::
-              {[:supervisor.child_spec | {module, any} | module], options :: keyword}
+              {[:supervisor.child_spec() | {module, any} | module], options :: keyword}
 
   @doc """
   Invoked by the terminator right before Broadway starts draining in-flight
@@ -27,16 +27,16 @@ defmodule Broadway.Producer do
   @optional_callbacks prepare_for_start: 2, 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)
+  def start_link(args, index, opts \\ []) do
+    GenStage.start_link(__MODULE__, {args, index}, opts)
   end
 
-  @spec push_messages(GenServer.server, [Message.t()]) :: :ok
+  @spec push_messages(GenServer.server(), [Message.t()]) :: :ok
   def push_messages(producer, messages) do
     GenStage.call(producer, {__MODULE__, :push_messages, messages})
   end
 
-  @spec drain(GenServer.server) :: :ok
+  @spec drain(GenServer.server()) :: :ok
   def drain(producer) do
     GenStage.cast(producer, {__MODULE__, :prepare_for_draining})
     GenStage.demand(producer, :accumulate)
@@ -44,10 +44,18 @@ defmodule Broadway.Producer do
   end
 
   @impl true
-  def init(args) do
+  def init({args, index}) do
     {module, arg} = args[:module]
     transformer = args[:transformer]
 
+    # Inject the topology index only if the args are a keyword list.
+    arg =
+      if Keyword.keyword?(arg) do
+        Keyword.put(arg, :broadway_index, index - 1)
+      else
+        arg
+      end
+
     state = %{
       module: module,
       module_state: nil,
diff --git a/lib/broadway/server.ex b/lib/broadway/server.ex
index 40767e2..170d785 100644
--- a/lib/broadway/server.ex
+++ b/lib/broadway/server.ex
@@ -103,23 +103,22 @@ defmodule Broadway.Server do
 
     n_producers = producer_config[:stages]
 
-    names =
+    names_and_specs =
       for index <- 1..n_producers do
-        producer_name(name_prefix(broadway_name), index)
-      end
-
-    specs =
-      for name <- names do
+        name = producer_name(name_prefix(broadway_name), index)
         opts = [name: name]
 
-        %{
-          start: {Producer, :start_link, [producer_config, opts]},
+        spec = %{
+          start: {Producer, :start_link, [producer_config, index, opts]},
           id: name,
           shutdown: shutdown
         }
+
+        {name, spec}
       end
 
-    {names, specs}
+    # We want to return {names, specs} here.
+    Enum.unzip(names_and_specs)
   end
 
   defp build_processors_specs(config, producers) do


Test patch:
diff --git a/test/broadway/producer_test.exs b/test/broadway/producer_test.exs
index ef544c9..c4fd197 100644
--- a/test/broadway/producer_test.exs
+++ b/test/broadway/producer_test.exs
@@ -71,7 +71,7 @@ defmodule Broadway.ProducerTest do
 
   test "init with bad return" do
     args = %{module: {ProducerWithBadReturn, []}}
-    assert Producer.init(args) == {:stop, {:bad_return_value, {:consumer, nil}}}
+    assert Producer.init({args, _index = 0}) == {:stop, {:bad_return_value, {:consumer, nil}}}
   end
 
   describe "wrap handle_demand" do
diff --git a/test/broadway_test.exs b/test/broadway_test.exs
index e906d00..e7bbf8f 100644
--- a/test/broadway_test.exs
+++ b/test/broadway_test.exs
@@ -289,6 +289,41 @@ defmodule BroadwayTest do
       assert [_, _, _, _, {Agent, _, _, _}] =
                Supervisor.which_children(Module.concat(broadway, "Broadway.Supervisor"))
     end
+
+    test "injects the :broadway_index option when the producer config is a kw list" do
+      defmodule ProducerWithTopologyIndex do
+        @behaviour Broadway.Producer
+
+        def init(opts) do
+          send(opts[:test_pid], {:init_called, opts})
+          {:producer, opts}
+        end
+
+        def handle_demand(_demand, state) do
+          {:noreply, [], state}
+        end
+      end
+
+      Broadway.start_link(Forwarder,
+        name: new_unique_name(),
+        producer: [module: {ProducerWithTopologyIndex, [test_pid: self()]}],
+        processors: [default: []]
+      )
+
+      assert_receive {:init_called, opts}
+      assert opts[:broadway_index] == 0
+
+      # If the producer config is not a kw list, the index is not injected.
+
+      Broadway.start_link(Forwarder,
+        name: new_unique_name(),
+        producer: [module: {ProducerWithTopologyIndex, %{test_pid: self()}}],
+        processors: [default: []]
+      )
+
+      assert_receive {:init_called, map}
+      refute Map.has_key?(map, :broadway_index)
+    end
   end
 
   test "push_messages/2" do