Repository: plataformatec/broadway
Pull request #55 (2019-03-08 03:09:44)

Allow pipelines without batchers

josevalim: Looks great! I have dropped some comments!
josevalim: Let's ship this!

Patch:
diff --git a/lib/broadway.ex b/lib/broadway.ex
index f527cde..1dfb80a 100644
--- a/lib/broadway.ex
+++ b/lib/broadway.ex
@@ -307,6 +307,8 @@ defmodule Broadway do
               context :: term
             ) :: [Message.t()]
 
+  @optional_callbacks handle_batch: 4
+
   @doc false
   defmacro __using__(opts) do
     quote location: :keep, bind_quoted: [opts: opts, module: __CALLER__.module] do
@@ -479,7 +481,7 @@ defmodule Broadway do
       context: [type: :any, default: :context_not_set],
       producers: [
         required: true,
-        type: :keyword_list,
+        type: :non_empty_keyword_list,
         keys: [
           *: [
             module: [required: true, type: :mod_arg],
@@ -490,7 +492,7 @@ defmodule Broadway do
       ],
       processors: [
         required: true,
-        type: :keyword_list,
+        type: :non_empty_keyword_list,
         keys: [
           *: [
             stages: [type: :pos_integer, default: System.schedulers_online() * 2],
@@ -500,7 +502,7 @@ defmodule Broadway do
         ]
       ],
       batchers: [
-        required: true,
+        default: [],
         type: :keyword_list,
         keys: [
           *: [
diff --git a/lib/broadway/options.ex b/lib/broadway/options.ex
index 6ee98ab..e9d6d68 100644
--- a/lib/broadway/options.ex
+++ b/lib/broadway/options.ex
@@ -1,7 +1,16 @@
 defmodule Broadway.Options do
   @moduledoc false
 
-  @types [:any, :keyword_list, :atom, :non_neg_integer, :pos_integer, :mfa, :mod_arg]
+  @types [
+    :any,
+    :keyword_list,
+    :non_empty_keyword_list,
+    :atom,
+    :non_neg_integer,
+    :pos_integer,
+    :mfa,
+    :mod_arg
+  ]
 
   def validate(opts, spec) do
     case validate_unknown_options(opts, spec) do
@@ -43,9 +52,9 @@ defmodule Broadway.Options do
   end
 
   defp validate_option(opts, key, spec) do
-    with {:ok, value} <- get_value(opts, key, spec),
-         :ok <- validate_type(spec[:type], key, value),
-         :ok <- validate_keys(spec[:keys], key, value) do
+    with {:ok, opts} <- validate_value(opts, key, spec),
+         value <- opts[key],
+         :ok <- validate_type(spec[:type], key, value) do
       if spec[:keys] do
         keys = normalize_keys(spec[:keys], value)
         validate(opts[key], keys)
@@ -55,14 +64,14 @@ defmodule Broadway.Options do
     end
   end
 
-  defp get_value(opts, key, spec) do
+  defp validate_value(opts, key, spec) do
     required? = Keyword.get(spec, :required, false)
     has_key? = Keyword.has_key?(opts, key)
     has_default? = Keyword.has_key?(spec, :default)
 
     case {required?, has_key?, has_default?} do
       {_, true, _} ->
-        {:ok, opts[key]}
+        {:ok, opts}
 
       {true, false, _} ->
         {:error,
@@ -70,21 +79,13 @@ defmodule Broadway.Options do
            inspect(Keyword.keys(opts))}
 
       {_, false, true} ->
-        {:ok, spec[:default]}
+        {:ok, Keyword.put(opts, key, spec[:default])}
 
       {_, false, false} ->
         :no_value
     end
   end
 
-  defp validate_keys(keys, key, value) do
-    if keys[:*] && value == [] do
-      {:error, "expected #{inspect(key)} to be a non-empty keyword list, got: #{inspect(value)}"}
-    else
-      :ok
-    end
-  end
-
   defp validate_type(:non_neg_integer, key, value) when not is_integer(value) or value < 0 do
     {:error, "expected #{inspect(key)} to be a non negative integer, got: #{inspect(value)}"}
   end
@@ -98,13 +99,21 @@ defmodule Broadway.Options do
   end
 
   defp validate_type(:keyword_list, key, value) do
-    if is_list(value) && Enum.all?(value, &tagged_tuple?/1) do
+    if keyword_list?(value) do
       :ok
     else
       {:error, "expected #{inspect(key)} to be a keyword list, got: #{inspect(value)}"}
     end
   end
 
+  defp validate_type(:non_empty_keyword_list, key, value) do
+    if keyword_list?(value) && value != [] do
+      :ok
+    else
+      {:error, "expected #{inspect(key)} to be a non-empty keyword list, got: #{inspect(value)}"}
+    end
+  end
+
   defp validate_type(:mfa, _key, {m, f, args}) when is_atom(m) and is_atom(f) and is_list(args) do
     :ok
   end
@@ -136,6 +145,10 @@ defmodule Broadway.Options do
   defp tagged_tuple?({key, _value}) when is_atom(key), do: true
   defp tagged_tuple?(_), do: false
 
+  defp keyword_list?(value) do
+    is_list(value) && Enum.all?(value, &tagged_tuple?/1)
+  end
+
   defp normalize_keys(keys, opts) do
     case keys[:*] do
       nil ->
diff --git a/lib/broadway/processor.ex b/lib/broadway/processor.ex
index d7348f6..a1f0c1a 100644
--- a/lib/broadway/processor.ex
+++ b/lib/broadway/processor.ex
@@ -34,14 +34,23 @@ defmodule Broadway.Processor do
   def handle_events(messages, _from, state) do
     {successful_messages, failed_messages} = handle_messages(messages, [], [], state)
 
+    {successful_messages_to_forward, successful_messages_to_ack} =
+      case state do
+        %{type: :consumer} ->
+          {[], successful_messages}
+
+        %{} ->
+          {successful_messages, []}
+      end
+
     try do
-      Acknowledger.ack_messages([], failed_messages)
+      Acknowledger.ack_messages(successful_messages_to_ack, failed_messages)
     catch
       kind, error ->
         Logger.error(Exception.format(kind, error, System.stacktrace()))
     end
 
-    {:noreply, successful_messages, state}
+    {:noreply, successful_messages_to_forward, state}
   end
 
   defp handle_messages([message | messages], successful, failed, state) do
@@ -73,6 +82,10 @@ defmodule Broadway.Processor do
     {Enum.reverse(successful), Enum.reverse(failed)}
   end
 
+  defp validate_message(%Message{} = message, []) do
+    message
+  end
+
   defp validate_message(%Message{batcher: batcher} = message, partitions) do
     if batcher not in partitions do
       raise "message was set to unknown batcher #{inspect(batcher)}. " <>
diff --git a/lib/broadway/server.ex b/lib/broadway/server.ex
index 246424b..be9fb6b 100644
--- a/lib/broadway/server.ex
+++ b/lib/broadway/server.ex
@@ -62,15 +62,12 @@ defmodule Broadway.Server do
     {producers_names, producers_specs} = build_producers_specs(config)
     {processors_names, processors_specs} = build_processors_specs(config, producers_names)
 
-    {consumers_names, batchers_consumers_specs} =
-      build_batchers_consumers_supervisors_specs(config, processors_names)
-
-    children = [
-      build_producer_supervisor_spec(config, producers_specs),
-      build_processor_supervisor_spec(config, processors_specs),
-      build_batcher_partition_supervisor_spec(config, batchers_consumers_specs),
-      build_terminator_spec(config, producers_names, processors_names, consumers_names)
-    ]
+    children =
+      [
+        build_producer_supervisor_spec(config, producers_specs),
+        build_processor_supervisor_spec(config, processors_specs)
+      ] ++
+        build_batcher_supervisor_and_terminator_specs(config, producers_names, processors_names)
 
     supervisor_opts = [
       name: :"#{config.name}.Supervisor",
@@ -155,10 +152,22 @@ defmodule Broadway.Server do
       end
 
     partitions = Keyword.keys(batchers_config)
-    dispatcher = {GenStage.PartitionDispatcher, partitions: partitions, hash: &{&1, &1.batcher}}
+
+    {type, dispatcher} =
+      case partitions do
+        [] ->
+          {:consumer, nil}
+
+        [_] ->
+          {:producer_consumer, GenStage.DemandDispatcher}
+
+        [_ | _] ->
+          {:producer_consumer,
+           {GenStage.PartitionDispatcher, partitions: partitions, hash: &{&1, &1.batcher}}}
+      end
 
     args = [
-      type: :producer_consumer,
+      type: type,
       resubscribe: resubscribe_interval,
       terminator: terminator,
       module: module,
@@ -183,6 +192,20 @@ defmodule Broadway.Server do
     {names, specs}
   end
 
+  defp build_batcher_supervisor_and_terminator_specs(config, producers_names, processors_names) do
+    if config[:batchers_config] == [] do
+      [build_terminator_spec(config, producers_names, processors_names, processors_names)]
+    else
+      {consumers_names, batchers_consumers_specs} =
+        build_batchers_consumers_supervisors_specs(config, processors_names)
+
+      [
+        build_batcher_partition_supervisor_spec(config, batchers_consumers_specs),
+        build_terminator_spec(config, producers_names, processors_names, consumers_names)
+      ]
+    end
+  end
+
   defp build_batchers_consumers_supervisors_specs(config, processors) do
     names_and_specs =
       for {key, _} = batcher_config <- config.batchers_config do
diff --git a/lib/broadway/subscriber.ex b/lib/broadway/subscriber.ex
index c3b9367..3a5227c 100644
--- a/lib/broadway/subscriber.ex
+++ b/lib/broadway/subscriber.ex
@@ -7,7 +7,7 @@ defmodule Broadway.Subscriber do
   # is set as one_for_all and max_restarts of 0 to the inner most
   # supervisor while the outer most is rest for one. This guarantees
   # that either all processess are running or none of them.
-  # 
+  #
   # For graceful shutdowns, we rely on cancellations with the help
   # of the terminator.
   @moduledoc false
@@ -42,7 +42,12 @@ defmodule Broadway.Subscriber do
       })
 
     Enum.each(names, &subscribe(&1, state))
-    {type, state, Keyword.take(options, [:dispatcher])}
+
+    if type == :consumer do
+      {type, state}
+    else
+      {type, state, Keyword.take(options, [:dispatcher])}
+    end
   end
 
   ## Callbacks


Test patch:
diff --git a/test/broadway/options_test.exs b/test/broadway/options_test.exs
index 9c563ab..12effa2 100644
--- a/test/broadway/options_test.exs
+++ b/test/broadway/options_test.exs
@@ -385,10 +385,10 @@ defmodule Broadway.OptionsTest do
                {:error, "expected :stages to be a positive integer, got: :an_atom"}
     end
 
-    test "validate empty keys" do
+    test "validate empty keys for :non_empty_keyword_list" do
       spec = [
         producers: [
-          type: :keyword_list,
+          type: :non_empty_keyword_list,
           keys: [
             *: [
               module: [required: true, type: :atom],
@@ -405,5 +405,44 @@ defmodule Broadway.OptionsTest do
       assert Options.validate(opts, spec) ==
                {:error, "expected :producers to be a non-empty keyword list, got: []"}
     end
+
+    test "allow empty keys for :keyword_list" do
+      spec = [
+        producers: [
+          type: :keyword_list,
+          keys: [
+            *: [
+              module: [required: true, type: :atom],
+              stages: [type: :pos_integer]
+            ]
+          ]
+        ]
+      ]
+
+      opts = [
+        producers: []
+      ]
+
+      assert Options.validate(opts, spec) == {:ok, opts}
+    end
+
+    test "default value for :keyword_list" do
+      spec = [
+        batchers: [
+          required: false,
+          default: [],
+          type: :keyword_list,
+          keys: [
+            *: [
+              stages: [type: :pos_integer, default: 1]
+            ]
+          ]
+        ]
+      ]
+
+      opts = []
+
+      assert Options.validate(opts, spec) == {:ok, [batchers: []]}
+    end
   end
 end
diff --git a/test/broadway_test.exs b/test/broadway_test.exs
index 06db57e..76e4976 100644
--- a/test/broadway_test.exs
+++ b/test/broadway_test.exs
@@ -88,6 +88,14 @@ defmodule BroadwayTest do
     end
   end
 
+  defmodule CustomHandlersWithoutHandleBatch do
+    use Broadway
+
+    def handle_message(_, message, %{handle_message: handler} = context) do
+      handler.(message, context)
+    end
+  end
+
   defmodule Transformer do
     def transform(event, test_pid: test_pid) do
       if event == :kill_producer do
@@ -392,6 +400,66 @@ defmodule BroadwayTest do
     end
   end
 
+  describe "pipeline without batchers" do
+    setup do
+      handle_message = fn message, _ ->
+        case message.data do
+          :fail -> Message.failed(message, "Failed message")
+          _ -> message
+        end
+      end
+
+      context = %{
+        handle_message: handle_message
+      }
+
+      broadway_name = new_unique_name()
+
+      {:ok, broadway} =
+        Broadway.start_link(CustomHandlersWithoutHandleBatch,
+          name: broadway_name,
+          context: context,
+          producers: [
+            default: [module: {ManualProducer, []}]
+          ],
+          processors: [default: [stages: 1, min_demand: 1, max_demand: 2]]
+        )
+
+      producer = get_producer(broadway_name, :default)
+
+      %{broadway_name: broadway_name, broadway: broadway, producer: producer}
+    end
+
+    test "no batcher supervisor is initialized", %{broadway_name: broadway_name} do
+      assert Process.whereis(:"#{broadway_name}.BatcherPartitionSupervisor") == nil
+    end
+
+    test "successful messages are marked as :ok", %{broadway: broadway} do
+      ref = Broadway.test_messages(broadway, [1, 2])
+      assert_receive {:ack, ^ref, [%{status: :ok}], []}
+      assert_receive {:ack, ^ref, [%{status: :ok}], []}
+    end
+
+    test "failed messages are marked as {:failed, reason}", %{broadway: broadway} do
+      ref = Broadway.test_messages(broadway, [:fail])
+      assert_receive {:ack, ^ref, _, [%{status: {:failed, "Failed message"}}]}
+    end
+
+    test "shutting down broadway waits until all events are processed",
+         %{broadway: broadway, producer: producer} do
+      # We suspend the producer to make sure that it doesn't process the messages early on
+      :sys.suspend(producer)
+      async_push_messages(producer, [1, 2, 3, 4])
+      Process.exit(broadway, :shutdown)
+      :sys.resume(producer)
+
+      assert_receive {:ack, _, [%{data: 1}], []}
+      assert_receive {:ack, _, [%{data: 2}], []}
+      assert_receive {:ack, _, [%{data: 3}], []}
+      assert_receive {:ack, _, [%{data: 4}], []}
+    end
+  end
+
   describe "put_batcher" do
     setup do
       broadway_name = new_unique_name()