add async flush, fix memory leak and race condition in client

Author: InvalidPathException

apps/kafkaesque_client/lib/kafkaesque_client/consumer.ex

@@ -25,7 +25,9 @@ defmodule KafkaesqueClient.Consumer do
     :stream_task,
     :consumer_group,
     :last_commit_time,
-    :metrics
+    :metrics,
+    :closed,
+    :auto_commit_timer
   ]
 
   @type t :: %__MODULE__{
@@ -40,7 +42,9 @@ defmodule KafkaesqueClient.Consumer do
           stream_task: Task.t() | nil,
           consumer_group: pid() | nil,
           last_commit_time: integer(),
-          metrics: map()
+          metrics: map(),
+          closed: boolean(),
+          auto_commit_timer: reference() | nil
         }
 
   # Client API
@@ -164,7 +168,9 @@ defmodule KafkaesqueClient.Consumer do
         polls: 0,
         commits: 0,
         errors: 0
-      }
+      },
+      closed: false,
+      auto_commit_timer: nil
     }
 
     # Start consumer group coordinator if we have a group
@@ -177,8 +183,11 @@ defmodule KafkaesqueClient.Consumer do
       end
 
     # Start auto-commit timer if enabled
-    if config.enable_auto_commit do
-      Process.send_after(self(), :auto_commit, config.auto_commit_interval_ms)
+    state = if config.enable_auto_commit do
+      timer_ref = Process.send_after(self(), :auto_commit, config.auto_commit_interval_ms)
+      %{state | auto_commit_timer: timer_ref}
+    else
+      state
     end
 
     {:ok, state}
@@ -272,6 +281,14 @@ defmodule KafkaesqueClient.Consumer do
   end
 
   def handle_call(:close, _from, state) do
+    # Mark as closed to prevent further operations
+    state = %{state | closed: true}
+
+    # Cancel auto-commit timer if active
+    if state.auto_commit_timer do
+      Process.cancel_timer(state.auto_commit_timer)
+    end
+
     # Commit offsets if auto-commit is enabled
     if state.config.enable_auto_commit do
       commit_current_offsets(state)
@@ -288,15 +305,20 @@ defmodule KafkaesqueClient.Consumer do
   end
 
   @impl true
+  def handle_info(:auto_commit, %{closed: true} = state) do
+    # Consumer is closed, ignore auto-commit
+    {:noreply, state}
+  end
+
   def handle_info(:auto_commit, state) do
     if should_auto_commit?(state) do
       commit_current_offsets(state)
     end
 
     # Schedule next auto-commit
-    Process.send_after(self(), :auto_commit, state.config.auto_commit_interval_ms)
+    timer_ref = Process.send_after(self(), :auto_commit, state.config.auto_commit_interval_ms)
 
-    {:noreply, %{state | last_commit_time: System.monotonic_time(:millisecond)}}
+    {:noreply, %{state | last_commit_time: System.monotonic_time(:millisecond), auto_commit_timer: timer_ref}}
   end
 
   def handle_info({:batch_received, batch}, state) do
@@ -531,6 +553,7 @@ defmodule KafkaesqueClient.Consumer do
 
   defp should_auto_commit?(state) do
     state.config.enable_auto_commit &&
+      not state.closed &&
       MapSet.size(state.subscriptions) > 0 &&
       map_size(state.position) > 0
   end

apps/kafkaesque_client/lib/kafkaesque_client/producer.ex

@@ -20,17 +20,19 @@ defmodule KafkaesqueClient.Producer do
     :batch_timer,
     :callbacks,
     :metrics,
-    :flush_waiters
+    :flush_waiters,
+    :cleanup_timer
   ]
 
   @type t :: %__MODULE__{
           config: map(),
           pool: atom() | pid(),
           batch: [ProducerRecord.t()],
           batch_timer: reference() | nil,
-          callbacks: %{reference() => function()},
+          callbacks: %{reference() => {function() | {:sync, GenServer.from()}, integer()}},
           metrics: map(),
-          flush_waiters: [GenServer.from()]
+          flush_waiters: [GenServer.from()],
+          cleanup_timer: reference() | nil
         }
 
   @type send_callback :: (
@@ -52,6 +54,7 @@ defmodule KafkaesqueClient.Producer do
   - `:linger_ms` - Time to wait before sending incomplete batches (default: 100)
   - `:max_retries` - Maximum number of retries (default: 3)
   - `:retry_backoff_ms` - Backoff between retries (default: 100)
+  - `:callback_timeout_ms` - Timeout for callbacks before cleanup (default: 30000)
   """
   def start_link(config) when is_map(config) do
     case Config.validate_producer_config(config) do
@@ -81,13 +84,22 @@ defmodule KafkaesqueClient.Producer do
   end
 
   @doc """
-  Flushes any pending records in the batch.
+  Flushes any pending records in the batch synchronously.
   """
   @spec flush(GenServer.server(), timeout()) :: :ok
   def flush(producer, timeout \\ 5_000) do
     GenServer.call(producer, :flush, timeout)
   end
 
+  @doc """
+  Flushes any pending records in the batch asynchronously.
+  Returns immediately without waiting for the batch to be sent.
+  """
+  @spec flush_async(GenServer.server()) :: :ok
+  def flush_async(producer) do
+    GenServer.cast(producer, :flush_async)
+  end
+
   @doc """
   Closes the producer, flushing any pending records.
   """
@@ -118,14 +130,20 @@ defmodule KafkaesqueClient.Producer do
         records_sent: 0,
         batches_sent: 0,
         errors: 0,
-        retries: 0
+        retries: 0,
+        callbacks_timeout: 0
       },
-      flush_waiters: []
+      flush_waiters: [],
+      cleanup_timer: nil
     }
 
     # Start batch timer if linger_ms > 0
     state = maybe_start_batch_timer(state)
 
+    # Start callback cleanup timer
+    cleanup_timer = Process.send_after(self(), :cleanup_callbacks, 5_000)
+    state = %{state | cleanup_timer: cleanup_timer}
+
     {:ok, state}
   end
 
@@ -135,7 +153,7 @@ defmodule KafkaesqueClient.Producer do
 
     new_callbacks =
       if callback do
-        Map.put(state.callbacks, ref, callback)
+        Map.put(state.callbacks, ref, {callback, System.monotonic_time(:millisecond)})
       else
         state.callbacks
       end
@@ -154,12 +172,20 @@ defmodule KafkaesqueClient.Producer do
     end
   end
 
+  def handle_cast(:flush_async, state) do
+    if Enum.empty?(state.batch) do
+      {:noreply, state}
+    else
+      {:noreply, send_batch(state)}
+    end
+  end
+
   @impl true
   def handle_call({:send_sync, record}, from, state) do
     ref = make_ref()
 
     # Store the caller to reply when batch is sent
-    new_callbacks = Map.put(state.callbacks, ref, {:sync, from})
+    new_callbacks = Map.put(state.callbacks, ref, {{:sync, from}, System.monotonic_time(:millisecond)})
 
     record_with_ref = {record, ref}
     new_batch = [record_with_ref | state.batch]
@@ -187,10 +213,13 @@ defmodule KafkaesqueClient.Producer do
     # Send any pending batch
     final_state = if Enum.empty?(state.batch), do: state, else: send_batch(state)
 
-    # Cancel timer if active
+    # Cancel timers if active
     if final_state.batch_timer do
       Process.cancel_timer(final_state.batch_timer)
     end
+    if final_state.cleanup_timer do
+      Process.cancel_timer(final_state.cleanup_timer)
+    end
 
     {:stop, :normal, :ok, final_state}
   end
@@ -204,6 +233,33 @@ defmodule KafkaesqueClient.Producer do
     {:noreply, send_batch(%{state | batch_timer: nil})}
   end
 
+  def handle_info(:cleanup_callbacks, state) do
+    now = System.monotonic_time(:millisecond)
+    timeout_ms = Map.get(state.config, :callback_timeout_ms, 30_000)
+
+    {expired, active} =
+      Map.split_with(state.callbacks, fn {_ref, {_cb, timestamp}} ->
+        now - timestamp > timeout_ms
+      end)
+
+    # Execute expired callbacks with timeout error
+    Enum.each(expired, fn {_ref, {callback, _timestamp}} ->
+      execute_callback_with_error(callback, :timeout)
+    end)
+
+    # Update metrics
+    new_metrics = if map_size(expired) > 0 do
+      %{state.metrics | callbacks_timeout: state.metrics.callbacks_timeout + map_size(expired)}
+    else
+      state.metrics
+    end
+
+    # Schedule next cleanup
+    cleanup_timer = Process.send_after(self(), :cleanup_callbacks, 5_000)
+
+    {:noreply, %{state | callbacks: active, cleanup_timer: cleanup_timer, metrics: new_metrics}}
+  end
+
   # Private Functions
 
   defp should_send_batch?(state) do
@@ -330,10 +386,10 @@ defmodule KafkaesqueClient.Producer do
       case Map.get(state.callbacks, ref) do
         nil -> :ok
 
-        {:sync, from} ->
+        {{:sync, from}, _timestamp} ->
           GenServer.reply(from, {:ok, metadata})
 
-        callback when is_function(callback, 1) ->
+        {callback, _timestamp} when is_function(callback, 1) ->
           Task.start(fn -> callback.({:ok, metadata}) end)
       end
     end)
@@ -345,15 +401,23 @@ defmodule KafkaesqueClient.Producer do
       case Map.get(state.callbacks, ref) do
         nil -> :ok
 
-        {:sync, from} ->
+        {{:sync, from}, _timestamp} ->
           GenServer.reply(from, {:error, reason})
 
-        callback when is_function(callback, 1) ->
+        {callback, _timestamp} when is_function(callback, 1) ->
           Task.start(fn -> callback.({:error, reason}) end)
       end
     end)
   end
 
+  defp execute_callback_with_error({:sync, from}, reason) do
+    GenServer.reply(from, {:error, reason})
+  end
+
+  defp execute_callback_with_error(callback, reason) when is_function(callback, 1) do
+    Task.start(fn -> callback.({:error, reason}) end)
+  end
+
   defp map_acks(:none), do: :ACKS_NONE
   defp map_acks(:leader), do: :ACKS_LEADER
   defp map_acks(:all), do: :ACKS_LEADER