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,57 @@ defmodule Broadway do
         }
         ```
 
+    * `[:broadway, :handle_failed, :start]` - Dispatched before the `c:handle_failed/2`
+      callback is invoked
+
+      * Measurement: `%{system_time: integer}`
+
+      * Metadata:
+
+        ```
+        %{
+          module: atom,
+          messages: [Broadway.Message.t],
+          context: term,
+          telemetry_span_context: reference
+        }
+        ```
+
+    * `[:broadway, :handle_failed, :stop]` - Dispatched after the `c:handle_failed/2`
+      callback has returned
+
+      * Measurement: `%{duration: native_time}`
+
+      * Metadata:
+
+        ```
+        %{
+          module: atom,
+          messages: [Broadway.Message.t],
+          context: term,
+          telemetry_span_context: reference
+        }
+        ```
+
+    * `[:broadway, :handle_failed, :exception]` - Dispatched if the `c:handle_failed/2`
+      callback raises an exception
+
+      * Measurement: `%{duration: native_time}`
+
+      * Metadata:
+
+        ```
+        %{
+          module: atom,
+          messages: [Broadway.Message.t],
+          context: term,
+          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

@@ -102,6 +102,19 @@ defmodule Broadway.Acknowledger do
   end
 
   defp handle_failed_messages(messages, module, context) do
+    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/processor_stage.ex

@@ -164,8 +164,7 @@ defmodule Broadway.Topology.ProcessorStage do
             index: state.partition,
             name: state.name,
             message: message,
-            context: state.context,
-            producer: state.producer
+            context: state.context
           },
           fn ->
             updated_message =
@@ -180,8 +179,7 @@ defmodule Broadway.Topology.ProcessorStage do
                index: state.partition,
                name: state.name,
                message: updated_message,
-               context: state.context,
-               producer: state.producer
+               context: state.context
              }}
           end
         )

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"]
@@ -72,7 +128,10 @@ defmodule Broadway.TelemetryTest do
       [:broadway, :processor, :message, :start],
       fn _event_name, _measurements, metadata, _config ->
         {_module, adapter_options} = metadata.producer
-        queue_name = adapter_options |> Keyword.fetch!(:queue_url) |> String.split("/") |> List.last()
+
+        queue_name =
+          adapter_options |> Keyword.fetch!(:queue_url) |> String.split("/") |> List.last()
+
         send(test_pid, {:queue_name, queue_name})
       end,
       %{}
@@ -92,4 +151,78 @@ 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.module == FailingMessageServer
+    assert [%Message{}] = start_metadata.messages
+
+    assert_receive {:telemetry_event, [:broadway, :handle_failed, :stop], stop_metadata}
+    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.module == FailingBatchServer
+    assert [%Message{}] = start_metadata.messages
+
+    assert_receive {:telemetry_event, [:broadway, :handle_failed, :stop], stop_metadata}
+    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