Add callback to allow instrumenting discarded count and function to get estimated buffer size (#257)

Author: mezz

lib/gen_stage.ex

@@ -397,7 +397,8 @@ defmodule GenStage do
   we will serve the existing demand, otherwise the event will be queued in
   `GenStage`'s internal buffer. In case events are being queued and not being
   consumed, a log message will be emitted when we exceed the `:buffer_size`
-  configuration.
+  configuration. This behavior can be customized by implementing the optional
+  `c:format_discarded/2` callback.
 
   While the implementation above is enough to solve the constraints above,
   a more robust implementation would have tighter control over the events
@@ -974,6 +975,15 @@ defmodule GenStage do
               | {:stop, reason, new_state}
             when new_state: term, reason: term
 
+  @doc """
+  Invoked when items are discarded from the buffer.
+
+  It receives the number of excess (discarded) items from this invocation.
+  This callback returns a boolean that controls whether the default error log for discarded items is printed or not.
+  Return true to print the log, return false to skip the log.
+  """
+  @callback format_discarded(discarded :: non_neg_integer, state :: term) :: boolean
+
   @doc """
   Invoked when a consumer is no longer subscribed to a producer.
 
@@ -1122,6 +1132,7 @@ defmodule GenStage do
     handle_cancel: 3,
     handle_demand: 2,
     handle_events: 3,
+    format_discarded: 2,
 
     # GenServer
     code_change: 3,
@@ -1693,6 +1704,14 @@ defmodule GenStage do
           "GenStage.stream/1 expects a list of subscriptions, got: #{inspect(subscriptions)}"
   end
 
+  @doc """
+  Returns the estimated number of buffered items for a producer.
+  """
+  @spec estimate_buffered_count(stage, timeout) :: non_neg_integer
+  def estimate_buffered_count(stage, timeout \\ 5000) do
+    call(stage, :"$estimate_buffered_count", timeout)
+  end
+
   ## Callbacks
 
   @compile :inline_list_funcs
@@ -1822,15 +1841,19 @@ defmodule GenStage do
     consumer_subscribe(current, to, opts, stage)
   end
 
+  def handle_call(:"$estimate_buffered_count", _from, stage) do
+    producer_estimate_buffered_count(stage)
+  end
+
   def handle_call(msg, from, %{mod: mod, state: state} = stage) do
     case mod.handle_call(msg, from, state) do
       {:reply, reply, events, state} when is_list(events) ->
-        stage = dispatch_events(events, length(events), stage)
-        {:reply, reply, %{stage | state: state}}
+        stage = dispatch_events(events, length(events), %{stage | state: state})
+        {:reply, reply, stage}
 
       {:reply, reply, events, state, :hibernate} when is_list(events) ->
-        stage = dispatch_events(events, length(events), stage)
-        {:reply, reply, %{stage | state: state}, :hibernate}
+        stage = dispatch_events(events, length(events), %{stage | state: state})
+        {:reply, reply, stage, :hibernate}
 
       {:stop, reason, reply, state} ->
         {:stop, reason, reply, %{stage | state: state}}
@@ -2106,12 +2129,12 @@ defmodule GenStage do
   defp handle_noreply_callback(return, stage) do
     case return do
       {:noreply, events, state} when is_list(events) ->
-        stage = dispatch_events(events, length(events), stage)
-        {:noreply, %{stage | state: state}}
+        stage = dispatch_events(events, length(events), %{stage | state: state})
+        {:noreply, stage}
 
       {:noreply, events, state, :hibernate} when is_list(events) ->
-        stage = dispatch_events(events, length(events), stage)
-        {:noreply, %{stage | state: state}, :hibernate}
+        stage = dispatch_events(events, length(events), %{stage | state: state})
+        {:noreply, stage, :hibernate}
 
       {:stop, reason, state} ->
         {:stop, reason, %{stage | state: state}}
@@ -2210,6 +2233,14 @@ defmodule GenStage do
     {:noreply, stage}
   end
 
+  defp maybe_format_discarded(mod, excess, state) do
+    if function_exported?(mod, :format_discarded, 2) do
+      mod.format_discarded(excess, state)
+    else
+      true
+    end
+  end
+
   defp producer_cancel(ref, kind, reason, stage) do
     %{consumers: consumers, monitors: monitors, state: state} = stage
 
@@ -2313,21 +2344,41 @@ defmodule GenStage do
     stage
   end
 
-  defp buffer_events(events, %{buffer: buffer, buffer_keep: keep} = stage) do
+  defp buffer_events(
+         events,
+         %{
+           mod: mod,
+           buffer: buffer,
+           buffer_keep: keep,
+           state: state
+         } = stage
+       ) do
     {buffer, excess, perms} = Buffer.store_temporary(buffer, events, keep)
 
     case excess do
       0 ->
         :ok
 
       excess ->
-        error_msg = 'GenStage producer ~tp has discarded ~tp events from buffer'
-        :error_logger.warning_msg(error_msg, [Utils.self_name(), excess])
+        if maybe_format_discarded(mod, excess, state) do
+          error_msg = 'GenStage producer ~tp has discarded ~tp events from buffer'
+          :error_logger.warning_msg(error_msg, [Utils.self_name(), excess])
+        end
     end
 
     :lists.foldl(&dispatch_info/2, %{stage | buffer: buffer}, perms)
   end
 
+  defp producer_estimate_buffered_count(%{type: :consumer} = stage) do
+    error_msg = 'Buffered count can only be requested for producers, GenStage ~tp is a consumer'
+    :error_logger.error_msg(error_msg, [Utils.self_name()])
+    {:reply, 0, stage}
+  end
+
+  defp producer_estimate_buffered_count(%{buffer: buffer} = stage) do
+    {:reply, Buffer.estimate_size(buffer), stage}
+  end
+
   ## Info helpers
 
   defp producer_info(msg, %{type: :consumer} = stage) do

test/gen_stage_test.exs

@@ -263,6 +263,43 @@ defmodule GenStageTest do
     end
   end
 
+  defmodule DiscardedBufferLogger do
+    @moduledoc """
+    Logs about any discarded items
+    """
+
+    use GenStage
+
+    def start_link(init, opts \\ []) do
+      GenStage.start_link(__MODULE__, init, opts)
+    end
+
+    def init(init) do
+      init
+    end
+
+    def sync_queue(stage, events) do
+      GenStage.call(stage, {:queue, events})
+    end
+
+    def handle_call({:queue, events}, _from, state) do
+      {:reply, :ok, events, state}
+    end
+
+    def format_discarded(discarded, %{log_discarded: log_discarded}) do
+      :error_logger.info_msg("DiscardedBufferLogger has discarded ~tp events from buffer", [
+        discarded
+      ])
+
+      log_discarded
+    end
+
+    def handle_demand(_demand, state) do
+      # We don't care about the demand
+      {:noreply, [], state}
+    end
+  end
+
   test "generates child_spec/1" do
     assert Counter.child_spec([:hello]) == %{
              id: Counter,
@@ -743,6 +780,93 @@ defmodule GenStageTest do
       assert_receive {:consumed, [:a, :b, :c, :d]}
       assert_receive {:consumed, [:e, :f, :g, :h]}
     end
+
+    test "calls optional format_discarded callback with discarded count when it exceeds configured size" do
+      {:ok, producer} =
+        DiscardedBufferLogger.start_link(
+          {:producer, %{log_discarded: false}, buffer_size: 5, buffer_keep: :first}
+        )
+
+      log =
+        capture_log(fn ->
+          DiscardedBufferLogger.sync_queue(producer, [:a, :b, :c, :d, :e, :f, :g, :h])
+        end)
+
+      assert log =~ "DiscardedBufferLogger has discarded 3 events from buffer"
+
+      log =
+        capture_log(fn ->
+          {:ok, consumer} = Forwarder.start_link({:consumer, self()})
+          :ok = GenStage.async_subscribe(consumer, to: producer, max_demand: 4, min_demand: 0)
+          assert_receive {:consumed, [:a, :b, :c, :d]}
+          assert_receive {:consumed, [:e]}
+        end)
+
+      assert log == ""
+    end
+
+    test "format_discarded can allow printing the default log when items are discarded" do
+      {:ok, producer} =
+        DiscardedBufferLogger.start_link(
+          {:producer, %{log_discarded: true}, buffer_size: 5, buffer_keep: :first}
+        )
+
+      log =
+        capture_log(fn ->
+          DiscardedBufferLogger.sync_queue(producer, [:a, :b, :c, :d, :e, :f, :g, :h])
+        end)
+
+      assert log =~ "GenStage producer #{inspect(producer)} has discarded 3 events from buffer"
+
+      log =
+        capture_log(fn ->
+          {:ok, consumer} = Forwarder.start_link({:consumer, self()})
+          :ok = GenStage.async_subscribe(consumer, to: producer, max_demand: 5, min_demand: 0)
+          assert_receive {:consumed, [:a, :b, :c, :d, :e]}
+        end)
+
+      assert log == ""
+    end
+
+    test "format_discarded can prevent printing the default log when items are discarded" do
+      {:ok, producer} =
+        DiscardedBufferLogger.start_link(
+          {:producer, %{log_discarded: false}, buffer_size: 5, buffer_keep: :first}
+        )
+
+      log =
+        capture_log(fn ->
+          DiscardedBufferLogger.sync_queue(producer, [:a, :b, :c, :d, :e, :f, :g, :h])
+        end)
+
+      assert not (log =~
+                    "GenStage producer #{inspect(producer)} has discarded 3 events from buffer")
+
+      log =
+        capture_log(fn ->
+          {:ok, consumer} = Forwarder.start_link({:consumer, self()})
+          :ok = GenStage.async_subscribe(consumer, to: producer, max_demand: 5, min_demand: 0)
+          assert_receive {:consumed, [:a, :b, :c, :d, :e]}
+        end)
+
+      assert log == ""
+    end
+
+    test "returns the correct buffer count when polled" do
+      {:ok, producer} = Counter.start_link({:producer, 0, buffer_size: :infinity})
+      0 = Counter.sync_queue(producer, [:a, :b, :c, :d, :e])
+      assert 5 == GenStage.estimate_buffered_count(producer)
+
+      0 = Counter.sync_queue(producer, [:f, :g, :h])
+      assert 8 == GenStage.estimate_buffered_count(producer)
+
+      {:ok, consumer} = Forwarder.start_link({:consumer, self()})
+      :ok = GenStage.async_subscribe(consumer, to: producer, max_demand: 4, min_demand: 0)
+      assert_receive {:consumed, [:a, :b, :c, :d]}
+      assert_receive {:consumed, [:e, :f, :g, :h]}
+
+      assert 0 == GenStage.estimate_buffered_count(producer)
+    end
   end
 
   describe "info" do