feat(telemetry): add :telemetry.span for handle_failed/2 callback

Signed-off-by: Yordis Prieto <yordis.prieto@gmail.com>

Author: yordis

lib/broadway.ex

@@ -798,6 +798,75 @@ defmodule Broadway do
         }
         ```
 
+    * `[:broadway, :handle_failed, :start]` - Dispatched before the `c:handle_failed/2`
+      callback is invoked
+
+      * Measurement: `%{system_time: integer}`
+
+      * Metadata:
+
+        ```
+        %{
+          topology_name: atom,
+          name: atom,
+          index: non_neg_integer,
+          processor_key: atom,           # only present when called from a processor
+          batch_info: Broadway.BatchInfo.t, # only present when called from a batch processor
+          module: atom,
+          messages: [Broadway.Message.t],
+          context: term,
+          producer: {atom, list},
+          telemetry_span_context: reference
+        }
+        ```
+
+    * `[:broadway, :handle_failed, :stop]` - Dispatched after the `c:handle_failed/2`
+      callback has returned
+
+      * Measurement: `%{duration: native_time}`
+
+      * Metadata:
+
+        ```
+        %{
+          topology_name: atom,
+          name: atom,
+          index: non_neg_integer,
+          processor_key: atom,           # only present when called from a processor
+          batch_info: Broadway.BatchInfo.t, # only present when called from a batch processor
+          module: atom,
+          messages: [Broadway.Message.t],
+          context: term,
+          producer: {atom, list},
+          telemetry_span_context: reference
+        }
+        ```
+
+    * `[:broadway, :handle_failed, :exception]` - Dispatched if the `c:handle_failed/2`
+      callback raises an exception
+
+      * Measurement: `%{duration: native_time}`
+
+      * Metadata:
+
+        ```
+        %{
+          topology_name: atom,
+          name: atom,
+          index: non_neg_integer,
+          processor_key: atom,           # only present when called from a processor
+          batch_info: Broadway.BatchInfo.t, # only present when called from a batch processor
+          module: atom,
+          messages: [Broadway.Message.t],
+          context: term,
+          producer: {atom, list},
+          kind: :error | :exit | :throw,
+          reason: term,
+          stacktrace: list,
+          telemetry_span_context: reference
+        }
+        ```
+
     * `[:broadway, :batcher, :start]` - Dispatched by a Broadway batcher before
       handling events
 

lib/broadway/acknowledger.ex

@@ -93,15 +93,28 @@ defmodule Broadway.Acknowledger do
   # Used by the processor and the batcher to maybe call c:handle_failed/2
   # on failed messages.
   @doc false
-  def maybe_handle_failed_messages(messages, module, context) do
+  def maybe_handle_failed_messages(messages, module, context, telemetry_metadata \\ %{}) do
     if function_exported?(module, :handle_failed, 2) and messages != [] do
-      handle_failed_messages(messages, module, context)
+      handle_failed_messages(messages, module, context, telemetry_metadata)
     else
       messages
     end
   end
 
-  defp handle_failed_messages(messages, module, context) do
+  defp handle_failed_messages(messages, module, context, telemetry_metadata) do
+    metadata = Map.merge(telemetry_metadata, %{module: module, messages: messages, context: context})
+
+    :telemetry.span(
+      [:broadway, :handle_failed],
+      metadata,
+      fn ->
+        result = do_handle_failed_messages(messages, module, context)
+        {result, Map.put(metadata, :messages, result)}
+      end
+    )
+  end
+
+  defp do_handle_failed_messages(messages, module, context) do
     module.handle_failed(messages, context)
   catch
     kind, reason ->

lib/broadway/topology/batch_processor_stage.ex

@@ -72,7 +72,14 @@ defmodule Broadway.Topology.BatchProcessorStage do
           Acknowledger.maybe_handle_failed_messages(
             failed_messages,
             state.module,
-            state.context
+            state.context,
+            %{
+              topology_name: state.topology_name,
+              name: state.name,
+              index: state.partition,
+              batch_info: batch_info,
+              producer: state.producer
+            }
           )
 
         if returned != size do

lib/broadway/topology/processor_stage.ex

@@ -91,7 +91,14 @@ defmodule Broadway.Topology.ProcessorStage do
           Acknowledger.maybe_handle_failed_messages(
             failed_messages,
             state.module,
-            state.context
+            state.context,
+            %{
+              topology_name: state.topology_name,
+              name: state.name,
+              index: state.partition,
+              processor_key: state.processor_key,
+              producer: state.producer
+            }
           )
 
         try do

test/broadway/telemetry_test.exs

@@ -1,6 +1,8 @@
 defmodule Broadway.TelemetryTest do
   use ExUnit.Case
 
+  alias Broadway.Message
+
   defmodule MessageServer do
     use Broadway
 
@@ -24,6 +26,60 @@ defmodule Broadway.TelemetryTest do
     end
   end
 
+  defmodule FailingMessageServer do
+    use Broadway
+
+    def start_link(opts) do
+      Broadway.start_link(__MODULE__,
+        name: opts[:name],
+        producer: [module: {Broadway.DummyProducer, []}],
+        processors: [default: []],
+        context: opts[:context]
+      )
+    end
+
+    @impl true
+    def handle_message(_, message, _context) do
+      Message.failed(message, "forced failure")
+    end
+
+    @impl true
+    def handle_failed(messages, %{test_pid: test_pid} = _context) do
+      send(test_pid, :handle_failed_called)
+      messages
+    end
+  end
+
+  defmodule FailingBatchServer do
+    use Broadway
+
+    def start_link(opts) do
+      Broadway.start_link(__MODULE__,
+        name: opts[:name],
+        producer: [module: {Broadway.DummyProducer, []}],
+        processors: [default: []],
+        batchers: [default: [batch_size: 1]],
+        context: opts[:context]
+      )
+    end
+
+    @impl true
+    def handle_message(_, message, _context) do
+      message
+    end
+
+    @impl true
+    def handle_batch(_batcher, messages, _batch_info, _context) do
+      Enum.map(messages, &Message.failed(&1, "forced batch failure"))
+    end
+
+    @impl true
+    def handle_failed(messages, %{test_pid: test_pid} = _context) do
+      send(test_pid, :handle_failed_called)
+      messages
+    end
+  end
+
   test "[:broadway, :processor, :message, :start] and :stop events include producer metadata" do
     test_pid = self()
     producer_options = [queue_url: "https://sqs.amazonaws.com/123456789/my-queue"]
@@ -92,4 +148,84 @@ defmodule Broadway.TelemetryTest do
 
     :telemetry.detach("test-extract-queue-name")
   end
+
+  test "[:broadway, :handle_failed, :start] and :stop are emitted from a processor" do
+    test_pid = self()
+
+    :telemetry.attach_many(
+      "test-handle-failed-processor",
+      [
+        [:broadway, :handle_failed, :start],
+        [:broadway, :handle_failed, :stop]
+      ],
+      fn event_name, _measurements, metadata, _config ->
+        send(test_pid, {:telemetry_event, event_name, metadata})
+      end,
+      %{}
+    )
+
+    {:ok, _pid} =
+      start_supervised(
+        {FailingMessageServer,
+         name: :test_handle_failed_processor, context: %{test_pid: test_pid}}
+      )
+
+    ref = Broadway.test_message(:test_handle_failed_processor, "data")
+
+    assert_receive {:telemetry_event, [:broadway, :handle_failed, :start], start_metadata}
+    assert start_metadata.topology_name == :test_handle_failed_processor
+    assert start_metadata.module == FailingMessageServer
+    assert start_metadata.processor_key == :default
+    assert [%Message{}] = start_metadata.messages
+
+    assert_receive {:telemetry_event, [:broadway, :handle_failed, :stop], stop_metadata}
+    assert stop_metadata.topology_name == :test_handle_failed_processor
+    assert stop_metadata.module == FailingMessageServer
+    assert [%Message{}] = stop_metadata.messages
+
+    assert_receive :handle_failed_called
+    assert_receive {:ack, ^ref, _successful, _failed}
+
+    :telemetry.detach("test-handle-failed-processor")
+  end
+
+  test "[:broadway, :handle_failed, :start] and :stop are emitted from a batch processor" do
+    test_pid = self()
+
+    :telemetry.attach_many(
+      "test-handle-failed-batch-processor",
+      [
+        [:broadway, :handle_failed, :start],
+        [:broadway, :handle_failed, :stop]
+      ],
+      fn event_name, _measurements, metadata, _config ->
+        send(test_pid, {:telemetry_event, event_name, metadata})
+      end,
+      %{}
+    )
+
+    {:ok, _pid} =
+      start_supervised(
+        {FailingBatchServer,
+         name: :test_handle_failed_batch_processor, context: %{test_pid: test_pid}}
+      )
+
+    ref = Broadway.test_message(:test_handle_failed_batch_processor, "data")
+
+    assert_receive {:telemetry_event, [:broadway, :handle_failed, :start], start_metadata}
+    assert start_metadata.topology_name == :test_handle_failed_batch_processor
+    assert start_metadata.module == FailingBatchServer
+    assert %Broadway.BatchInfo{} = start_metadata.batch_info
+    assert [%Message{}] = start_metadata.messages
+
+    assert_receive {:telemetry_event, [:broadway, :handle_failed, :stop], stop_metadata}
+    assert stop_metadata.topology_name == :test_handle_failed_batch_processor
+    assert stop_metadata.module == FailingBatchServer
+    assert [%Message{}] = stop_metadata.messages
+
+    assert_receive :handle_failed_called
+    assert_receive {:ack, ^ref, _successful, _failed}
+
+    :telemetry.detach("test-handle-failed-batch-processor")
+  end
 end