Extract PipelineUnit and move configs to separated structs

Author: Gleb Ivanov

.dialyzer_ignore.exs

@@ -1,4 +1,4 @@
 [
   # If we compile with another @storage_impl lib/kafka_batcher/temp_storage.ex:33 become reachable
-  {"lib/kafka_batcher/temp_storage.ex", :guard_fail, 30}
+  {"lib/kafka_batcher/temp_storage.ex", :guard_fail, 32}
 ]

lib/kafka_batcher.ex

@@ -67,6 +67,8 @@ defmodule KafkaBatcher do
     """
     defstruct key: "", value: "", headers: []
 
-    @type t :: %MessageObject{key: binary(), value: map() | binary(), headers: list()}
+    @type key :: binary()
+    @type value :: map() | binary()
+    @type t :: %MessageObject{key: key(), value: value(), headers: list()}
   end
 end

lib/kafka_batcher/accumulator.ex

@@ -5,7 +5,15 @@ defmodule KafkaBatcher.Accumulator do
   See details how it works in KafkaBatcher.Accumulator.State module
   """
 
-  alias KafkaBatcher.{Accumulator.State, MessageObject, TempStorage}
+  alias KafkaBatcher.{
+    Accumulator,
+    Accumulator.State,
+    MessageObject,
+    PipelineUnit,
+    Producers,
+    TempStorage
+  }
+
   alias KafkaBatcher.Behaviours.Collector, as: CollectorBehaviour
 
   @error_notifier Application.compile_env(:kafka_batcher, :error_notifier, KafkaBatcher.DefaultErrorNotifier)
@@ -14,25 +22,34 @@ defmodule KafkaBatcher.Accumulator do
   use GenServer
   require Logger
 
-  def start_link(args) do
-    GenServer.start_link(__MODULE__, args, name: reg_name(args))
+  @spec start_link(PipelineUnit.t()) :: GenServer.on_start()
+  def start_link(%PipelineUnit{} = pipeline_unit) do
+    GenServer.start_link(
+      __MODULE__,
+      pipeline_unit,
+      name: reg_name(pipeline_unit)
+    )
   end
 
   @doc "Returns a specification to start this module under a supervisor"
-  def child_spec(args) do
-    {accumulator_mod, args} = Keyword.pop(args, :accumulator_mod, __MODULE__)
-
+  @spec child_spec(PipelineUnit.t()) :: Supervisor.child_spec()
+  def child_spec(%PipelineUnit{} = pipeline_unit) do
     %{
-      id: reg_name(args),
-      start: {accumulator_mod, :start_link, [args]}
+      id: reg_name(pipeline_unit),
+      start: {
+        PipelineUnit.get_accumulator_mod(pipeline_unit),
+        :start_link,
+        [pipeline_unit]
+      }
     }
   end
 
   @doc """
   Finds appropriate Accumulator process by topic & partition and dispatches `event` to it
   """
-  def add_event(%MessageObject{} = event, topic_name, partition \\ nil) do
-    GenServer.call(reg_name(topic_name: topic_name, partition: partition), {:add_event, event})
+  @spec add_event(MessageObject.t(), PipelineUnit.t()) :: :ok | {:error, term()}
+  def add_event(%MessageObject{} = event, %PipelineUnit{} = pipeline_unit) do
+    GenServer.call(reg_name(pipeline_unit), {:add_event, event})
   catch
     _, _reason ->
       Logger.warning("KafkaBatcher: Couldn't get through to accumulator")
@@ -43,15 +60,17 @@ defmodule KafkaBatcher.Accumulator do
   ## Callbacks
   ##
   @impl GenServer
-  def init(args) do
+  def init(%PipelineUnit{} = pipeline_unit) do
     Process.flag(:trap_exit, true)
-    state = build_state(args)
+
+    topic_name = PipelineUnit.get_topic_name(pipeline_unit)
+    partition = PipelineUnit.get_partition(pipeline_unit)
 
     Logger.debug("""
-      KafkaBatcher: Accumulator process started: topic #{state.topic_name} partition #{state.partition} pid #{inspect(self())}
+      KafkaBatcher: Accumulator process started: topic #{topic_name} partition #{partition} pid #{inspect(self())}
     """)
 
-    {:ok, state}
+    {:ok, %State{pipeline_unit: pipeline_unit}}
   end
 
   @impl GenServer
@@ -78,7 +97,7 @@ defmodule KafkaBatcher.Accumulator do
         {:noreply, new_state}
 
       {:error, _reason, new_state} ->
-        state.collector.set_lock()
+        PipelineUnit.get_collector(state.pipeline_unit).set_lock()
         {:noreply, new_state}
     end
   end
@@ -92,7 +111,7 @@ defmodule KafkaBatcher.Accumulator do
   def handle_info(term, state) do
     Logger.warning("""
       KafkaBatcher: Unknown message #{inspect(term)} to #{__MODULE__}.handle_info/2.
-      Current state: #{inspect(drop_sensitive(state))}
+      Current state: #{inspect(state)}
     """)
 
     {:noreply, state}
@@ -104,12 +123,14 @@ defmodule KafkaBatcher.Accumulator do
   end
 
   @impl GenServer
-  def format_status(_reason, [pdict, state]) do
-    [pdict, drop_sensitive(state)]
-  end
-
-  defp drop_sensitive(%State{config: config} = state) do
-    %State{state | config: Keyword.drop(config, [:sasl])}
+  def format_status(_reason, [pdict, %State{} = state]) do
+    [
+      pdict,
+      %State{
+        state
+        | pipeline_unit: PipelineUnit.drop_sensitive(state.pipeline_unit)
+      }
+    ]
   end
 
   defp cleanup(%{pending_messages: [], messages_to_produce: []}) do
@@ -122,7 +143,11 @@ defmodule KafkaBatcher.Accumulator do
   end
 
   defp set_cleanup_timer_if_not_exists(%State{cleanup_timer_ref: nil} = state) do
-    ref = :erlang.start_timer(state.max_wait_time, self(), :cleanup)
+    %PipelineUnit{
+      accumulator_config: %Accumulator.Config{max_wait_time: max_wait_time}
+    } = state.pipeline_unit
+
+    ref = :erlang.start_timer(max_wait_time, self(), :cleanup)
     %State{state | cleanup_timer_ref: ref}
   end
 
@@ -161,37 +186,26 @@ defmodule KafkaBatcher.Accumulator do
 
   @spec produce_list(messages :: [CollectorBehaviour.event()], state :: State.t()) :: :ok | {:error, any()}
   defp produce_list(messages, state) when is_list(messages) do
-    @producer.produce_list(messages, state.topic_name, state.partition, state.config)
+    @producer.produce_list(state.config, messages, state.topic_name, state.partition)
   catch
     _, reason ->
       {:error, reason}
   end
 
-  defp build_state(args) do
-    config = Keyword.fetch!(args, :config)
-
-    %State{
-      topic_name: Keyword.fetch!(args, :topic_name),
-      partition: Keyword.get(args, :partition),
-      config: config,
-      batch_flusher: Keyword.fetch!(config, :batch_flusher),
-      batch_size: Keyword.fetch!(config, :batch_size),
-      max_wait_time: Keyword.fetch!(config, :max_wait_time),
-      min_delay: Keyword.fetch!(config, :min_delay),
-      max_batch_bytesize: Keyword.fetch!(config, :max_batch_bytesize),
-      collector: Keyword.fetch!(args, :collector)
-    }
-  end
+  defp reg_name(%PipelineUnit{} = pipeline_unit) do
+    %PipelineUnit{
+      producer_config: %Producers.Config{client_name: client_name}
+    } = pipeline_unit
 
-  defp reg_name(args) do
-    topic_name = Keyword.fetch!(args, :topic_name)
+    topic_name = PipelineUnit.get_topic_name(pipeline_unit)
+    partition = PipelineUnit.get_partition(pipeline_unit)
 
-    case Keyword.get(args, :partition) do
+    case partition do
       nil ->
-        :"#{__MODULE__}.#{topic_name}"
+        :"#{__MODULE__}.#{client_name}.#{topic_name}"
 
       partition ->
-        :"#{__MODULE__}.#{topic_name}.#{partition}"
+        :"#{__MODULE__}.#{client_name}.#{topic_name}.#{partition}"
     end
   end
 end

lib/kafka_batcher/accumulator/config.ex

@@ -0,0 +1,62 @@
+defmodule KafkaBatcher.Accumulator.Config do
+  @moduledoc false
+
+  @type t :: %__MODULE__{
+          collector: module(),
+          topic_name: String.t(),
+          partition: pos_integer() | nil,
+          batch_flusher: module(),
+          max_wait_time: pos_integer(),
+          batch_size: pos_integer(),
+          min_delay: non_neg_integer(),
+          max_batch_bytesize: pos_integer(),
+          max_accumulator_restarts: pos_integer(),
+          accumulator_mod: module()
+        }
+
+  @enforce_keys [:collector, :topic_name]
+  defstruct @enforce_keys ++
+              [
+                :partition,
+                batch_flusher: KafkaBatcher.Accumulator.DefaultBatchFlusher,
+                max_wait_time: 1_000,
+                batch_size: 10,
+                min_delay: 0,
+                max_batch_bytesize: 1_000_000,
+                max_accumulator_restarts: 100,
+                accumulator_mod: KafkaBatcher.Accumulator
+              ]
+
+  @spec to_kwlist(t()) :: Keyword.t()
+  def to_kwlist(%__MODULE__{} = config) do
+    [
+      collector: config.collector,
+      topic_name: config.topic_name,
+      partition: config.partition,
+      batch_flusher: config.batch_flusher,
+      max_wait_time: config.max_wait_time,
+      batch_size: config.batch_size,
+      min_delay: config.min_delay,
+      max_batch_bytesize: config.max_batch_bytesize,
+      max_accumulator_restarts: config.max_accumulator_restarts,
+      accumulator_mod: config.accumulator_mod
+    ]
+  end
+
+  @spec build!(opts :: Keyword.t()) :: t()
+  def build!(opts) do
+    opts
+    |> Keyword.take([
+      :collector,
+      :topic_name,
+      :batch_flusher,
+      :max_wait_time,
+      :batch_size,
+      :min_delay,
+      :max_batch_bytesize,
+      :max_accumulator_restarts,
+      :accumulator_mod
+    ])
+    |> then(&struct!(__MODULE__, &1))
+  end
+end

lib/kafka_batcher/accumulator/state.ex

@@ -11,46 +11,37 @@ defmodule KafkaBatcher.Accumulator.State do
     * timer expired (in case when a few events arrived timer helps to control that the max waiting time is not exceeded)
   """
 
-  alias KafkaBatcher.{Accumulator.State, MessageObject}
+  alias KafkaBatcher.{
+    Accumulator,
+    Accumulator.State,
+    MessageObject,
+    PipelineUnit
+  }
+
   @error_notifier Application.compile_env(:kafka_batcher, :error_notifier, KafkaBatcher.DefaultErrorNotifier)
 
   @type t :: %State{
-          topic_name: binary(),
-          partition: non_neg_integer() | nil,
-          config: Keyword.t(),
+          pipeline_unit: KafkaBatcher.PipelineUnit.t(),
           pending_messages: list(),
           last_produced_at: non_neg_integer(),
-          batch_flusher: atom(),
-          batch_size: non_neg_integer(),
-          max_wait_time: non_neg_integer(),
-          min_delay: non_neg_integer(),
-          max_batch_bytesize: non_neg_integer(),
           batch_bytesize: non_neg_integer(),
           pending_messages_count: non_neg_integer(),
-          producer_config: Keyword.t(),
           messages_to_produce: list(),
           cleanup_timer_ref: reference() | nil,
-          status: atom(),
-          collector: atom() | nil
+          status: atom()
         }
 
-  defstruct topic_name: nil,
-            partition: nil,
-            config: [],
-            pending_messages: [],
-            last_produced_at: 0,
-            batch_flusher: KafkaBatcher.Accumulator.DefaultBatchFlusher,
-            batch_size: 0,
-            max_wait_time: 0,
-            min_delay: 0,
-            max_batch_bytesize: 0,
-            batch_bytesize: 0,
-            pending_messages_count: 0,
-            producer_config: [],
-            messages_to_produce: [],
-            cleanup_timer_ref: nil,
-            status: :continue,
-            collector: nil
+  @enforce_keys [:pipeline_unit]
+  defstruct @enforce_keys ++
+              [
+                pending_messages: [],
+                last_produced_at: 0,
+                batch_bytesize: 0,
+                pending_messages_count: 0,
+                messages_to_produce: [],
+                cleanup_timer_ref: nil,
+                status: :continue
+              ]
 
   @spec add_new_message(State.t(), MessageObject.t(), non_neg_integer()) :: State.t()
   def add_new_message(%State{} = state, %MessageObject{key: key, value: value} = event, now) do
@@ -88,15 +79,22 @@ defmodule KafkaBatcher.Accumulator.State do
   end
 
   defp consider_max_bytesize(%State{status: :continue, batch_bytesize: batch_bytesize} = state, new_message) do
+    %PipelineUnit{
+      accumulator_config: %Accumulator.Config{max_batch_bytesize: max_batch_bytesize}
+    } = state.pipeline_unit
+
+    topic_name = PipelineUnit.get_topic_name(state.pipeline_unit)
+    partition = PipelineUnit.get_partition(state.pipeline_unit)
+
     message_size = :erlang.external_size(new_message)
 
-    case batch_bytesize + message_size >= state.max_batch_bytesize do
-      true when message_size >= state.max_batch_bytesize ->
+    case batch_bytesize + message_size >= max_batch_bytesize do
+      true when message_size >= max_batch_bytesize ->
         @error_notifier.report(
           type: "KafkaBatcherProducerError",
           message: """
-          event#produce topic=#{state.topic_name} partition=#{state.partition}.
-          Message size #{inspect(message_size)} exceeds limit #{inspect(state.max_batch_bytesize)}
+          event#produce topic=#{topic_name} partition=#{partition}.
+          Message size #{inspect(message_size)} exceeds limit #{inspect(max_batch_bytesize)}
           """
         )
 
@@ -111,7 +109,11 @@ defmodule KafkaBatcher.Accumulator.State do
   end
 
   defp consider_max_size_and_wait_time(%State{status: :continue} = state, now) do
-    if state.pending_messages_count >= state.batch_size and now - state.last_produced_at >= state.min_delay do
+    %PipelineUnit{
+      accumulator_config: %Accumulator.Config{batch_size: batch_size, min_delay: min_delay}
+    } = state.pipeline_unit
+
+    if state.pending_messages_count >= batch_size and now - state.last_produced_at >= min_delay do
       mark_as_ready(state)
     else
       state
@@ -121,7 +123,11 @@ defmodule KafkaBatcher.Accumulator.State do
   defp consider_max_size_and_wait_time(%State{status: :ready} = state, _), do: state
 
   defp consider_istant_flush(%State{status: :continue} = state, key, value) do
-    if state.batch_flusher.flush?(key, value) do
+    %PipelineUnit{
+      accumulator_config: %Accumulator.Config{batch_flusher: batch_flusher}
+    } = state.pipeline_unit
+
+    if batch_flusher.flush?(key, value) do
       mark_as_ready(state)
     else
       state