feat: implement shared client

lib/broadway_kafka/brod_client.ex

@@ -47,6 +47,8 @@ defmodule BroadwayKafka.BrodClient do
 
   @default_begin_offset :assigned
 
+  @default_shared_client false
+
   @impl true
   def init(opts) do
     with {:ok, hosts} <- validate(opts, :hosts, required: true),
@@ -62,6 +64,8 @@ defmodule BroadwayKafka.BrodClient do
            validate(opts, :offset_reset_policy, default: @default_offset_reset_policy),
          {:ok, begin_offset} <-
            validate(opts, :begin_offset, default: @default_begin_offset),
+         {:ok, shared_client} <-
+           validate(opts, :shared_client, default: @default_shared_client),
          {:ok, group_config} <- validate_group_config(opts),
          {:ok, fetch_config} <- validate_fetch_config(opts),
          {:ok, client_config} <- validate_client_config(opts) do
@@ -77,14 +81,16 @@ defmodule BroadwayKafka.BrodClient do
          begin_offset: begin_offset,
          group_config: [{:offset_commit_policy, @offset_commit_policy} | group_config],
          fetch_config: Map.new(fetch_config || []),
-         client_config: client_config
+         client_config: client_config,
+         shared_client: shared_client,
+         shared_client_id: build_shared_client_id(opts)
        }}
     end
   end
 
   @impl true
   def setup(stage_pid, client_id, callback_module, config) do
-    with :ok <- :brod.start_client(config.hosts, client_id, config.client_config),
+    with :ok <- do_start_brod_client(config.hosts, client_id, config.client_config),
          {:ok, group_coordinator} <-
            start_link_group_coordinator(stage_pid, client_id, callback_module, config) do
       Process.monitor(client_id)
@@ -147,6 +153,18 @@ defmodule BroadwayKafka.BrodClient do
     end
   end
 
+  @impl true
+  def shared_client_child_spec(config) do
+    [
+      %{
+        id: config.shared_client_id,
+        start:
+          {:brod, :start_link_client,
+           [config.hosts, config.shared_client_id, config.client_config]}
+      }
+    ]
+  end
+
   defp lookup_offset(hosts, topic, partition, policy, client_config) do
     case :brod.resolve_offset(hosts, topic, partition, policy, client_config) do
       {:ok, offset} ->
@@ -268,6 +286,9 @@ defmodule BroadwayKafka.BrodClient do
   defp validate_option(:client_id_prefix, value) when not is_binary(value),
     do: validation_error(:client_id_prefix, "a string", value)
 
+  defp validate_option(:shared_client, value) when not is_boolean(value),
+    do: validation_error(:shared_client, "a boolean", value)
+
   defp validate_option(:sasl, :undefined),
     do: {:ok, :undefined}
 
@@ -387,4 +408,29 @@ defmodule BroadwayKafka.BrodClient do
   end
 
   defp parse_hosts(hosts), do: hosts
+
+  defp build_shared_client_id(opts) do
+    if opts[:shared_client] do
+      prefix = get_in(opts, [:client_config, :client_id_prefix])
+      broadway_name = opts[:broadway][:name]
+      :"#{prefix}#{Module.concat(broadway_name, SharedClient)}"
+    end
+  end
+
+  defp do_start_brod_client(hosts, client_id, client_config) do
+    case :brod.start_client(hosts, client_id, client_config) do
+      :ok ->
+        :ok
+
+      # Because  we are starting the client on the broadway supervison tree
+      # instead of the :brod supervisor, the already_started error
+      # is not properly handled by :brod.start_client/3 for shared clients
+      # So we must handle it here.
+      {:error, {{:already_started, _}, _}} ->
+        :ok
+
+      error ->
+        error
+    end
+  end
 end

lib/broadway_kafka/kafka_client.ex

@@ -9,7 +9,9 @@ defmodule BroadwayKafka.KafkaClient do
            offset_commit_on_ack: boolean,
            topics: [:brod.topic()],
            group_config: keyword,
-           client_config: keyword
+           client_config: keyword,
+           shared_client: boolean(),
+           shared_client_id: atom() | nil
          }
 
   @typep offset_reset_policy :: :earliest | :latest
@@ -52,4 +54,9 @@ defmodule BroadwayKafka.KafkaClient do
   @callback update_topics(:brod.group_coordinator(), [:brod.topic()]) :: :ok
   @callback connected?(:brod.client()) :: boolean
   @callback disconnect(:brod.client()) :: :ok
+
+  @callback shared_client_child_spec(config()) :: [child_spec]
+            when child_spec: :supervisor.child_spec() | {module, any} | module
+
+  @optional_callbacks shared_client_child_spec: 1
 end

lib/broadway_kafka/producer.ex

@@ -49,6 +49,10 @@ defmodule BroadwayKafka.Producer do
       When set to `:reset`, the starting offset will be dictated by the `:offset_reset_policy` option, either
       starting from the `:earliest` or the `:latest` offsets of the topic. Default is `:assigned`.
 
+    * `:shared_client` - Optional. A boolean that defines how many clients will be started.
+    If `true`, only one shared client will be started for all producers, if `false` each producer
+    will have it's own client. Default is `false`
+
     * `:group_config` - Optional. A list of options used to configure the group
       coordinator. See the ["Group config options"](#module-group-config-options) section below for a list of all available
       options.
@@ -233,7 +237,7 @@ defmodule BroadwayKafka.Producer do
 
     client = opts[:client] || BroadwayKafka.BrodClient
 
-    case client.init(opts) do
+    case opts[:initialized_client_config] || client.init(opts) do
       {:error, message} ->
         raise ArgumentError, "invalid options given to #{inspect(client)}.init/1, " <> message
 
@@ -246,7 +250,9 @@ defmodule BroadwayKafka.Producer do
           |> drain_after_revoke_table_init!()
 
         prefix = get_in(config, [:client_config, :client_id_prefix])
-        client_id = :"#{prefix}#{Module.concat([producer_name, Client])}"
+
+        client_id =
+          config[:shared_client_id] || :"#{prefix}#{Module.concat([producer_name, Client])}"
 
         max_demand =
           with [{_first, processor_opts}] <- opts[:broadway][:processors],
@@ -271,7 +277,8 @@ defmodule BroadwayKafka.Producer do
           demand: 0,
           shutting_down?: false,
           buffer: :queue.new(),
-          max_demand: max_demand
+          max_demand: max_demand,
+          shared_client: config.shared_client
         }
 
         {:producer, connect(state)}
@@ -509,7 +516,29 @@ defmodule BroadwayKafka.Producer do
       |> Keyword.put(:processors, [updated_processor_entry | other_processors_entries])
       |> Keyword.put(:batchers, updated_batchers_entries)
 
-    {allocators, updated_opts}
+    {producer_mod, producer_opts} = opts[:producer][:module]
+
+    {extra_child_specs, initialized_client_config} =
+      if producer_opts[:shared_client] do
+        client = producer_opts[:client] || BroadwayKafka.BrodClient
+
+        case client.init(Keyword.put(producer_opts, :broadway, opts)) do
+          {:error, message} ->
+            raise ArgumentError, "invalid options given to #{client}.init/1, " <> message
+
+          {:ok, config} = result ->
+            {client.shared_client_child_spec(config), result}
+        end
+      else
+        {[], nil}
+      end
+
+    new_producer_opts =
+      Keyword.put(producer_opts, :initialized_client_config, initialized_client_config)
+
+    updated_opts = put_in(updated_opts, [:producer, :module], {producer_mod, new_producer_opts})
+
+    {allocators ++ extra_child_specs, updated_opts}
   end
 
   @impl :brod_group_member
@@ -547,7 +576,11 @@ defmodule BroadwayKafka.Producer do
   def terminate(_reason, state) do
     %{client: client, group_coordinator: group_coordinator, client_id: client_id} = state
     group_coordinator && Process.exit(group_coordinator, :shutdown)
-    client.disconnect(client_id)
+
+    if state.shared_client == false do
+      client.disconnect(client_id)
+    end
+
     :ok
   end
 

test/brod_client_test.exs

@@ -363,6 +363,77 @@ defmodule BroadwayKafka.BrodClientTest do
 
       assert {:ok, %{client_config: [query_api_versions: false]}} = BrodClient.init(opts)
     end
+
+    test ":shared_client is an optional boolean" do
+      opts = Keyword.put(@opts, :shared_client, "true")
+
+      assert BrodClient.init(opts) ==
+               {:error, "expected :shared_client to be a boolean, got: \"true\""}
+
+      opts =
+        @opts
+        |> Keyword.put(:shared_client, true)
+        |> Keyword.put(:broadway, name: :my_broadway_name)
+        |> put_in([:client_config, :client_id_prefix], "my_prefix.")
+
+      assert {:ok, %{shared_client: true}} = BrodClient.init(opts)
+    end
+
+    test "return shared_client_id when :shared_client is true" do
+      opts =
+        @opts
+        |> Keyword.put(:shared_client, true)
+        |> Keyword.put(:broadway, name: :my_broadway_name)
+        |> put_in([:client_config, :client_id_prefix], "my_prefix.")
+
+      assert {:ok,
+              %{
+                shared_client: true,
+                shared_client_id: :"my_prefix.Elixir.my_broadway_name.SharedClient"
+              }} =
+               BrodClient.init(opts)
+
+      opts =
+        @opts
+        |> Keyword.put(:shared_client, false)
+        |> Keyword.put(:broadway, name: :my_broadway_name)
+        |> put_in([:client_config, :client_id_prefix], "my_prefix.")
+
+      assert {:ok,
+              %{
+                shared_client: false,
+                shared_client_id: nil
+              }} =
+               BrodClient.init(opts)
+    end
+  end
+
+  describe "shared_client_child_spec" do
+    test "should return child spec" do
+      module_opts =
+        @opts
+        |> Keyword.put(:shared_client, true)
+        |> Keyword.put(:client_config, client_id_prefix: "my_prefix.")
+
+      broadway_opts = [
+        name: :my_broadway
+      ]
+
+      {:ok, config} = BrodClient.init(Keyword.put(module_opts, :broadway, broadway_opts))
+
+      assert child_specs = BrodClient.shared_client_child_spec(config)
+
+      assert [
+               %{
+                 id: shared_client_id,
+                 start: {:brod, :start_link_client, [hosts, shared_client_id, client_config]}
+               }
+             ] = child_specs
+
+      assert [{:host, 9092}] = hosts
+      assert :"my_prefix.Elixir.my_broadway.SharedClient" = shared_client_id
+      assert [client_id_prefix: "my_prefix."] = client_config
+    end
   end
 
   defmodule FakeSaslMechanismPlugin do