Instrument batching in opentelemetry_broadway

instrumentation/opentelemetry_broadway/lib/opentelemetry_broadway.ex

@@ -53,6 +53,22 @@ defmodule OpentelemetryBroadway do
  []
  )
 
+ :ok =
+ :telemetry.attach(
+ "#{__MODULE__}.batch_start",
+ [:broadway, :batch_processor, :start],
+ &__MODULE__.handle_batch_start/4,
+ []
+ )
+
+ :ok =
+ :telemetry.attach(
+ "#{__MODULE__}.batch_stop",
+ [:broadway, :batch_processor, :stop],
+ &__MODULE__.handle_batch_stop/4,
+ []
+ )
+
  :ok
  end
 
@@ -137,6 +153,58 @@ defmodule OpentelemetryBroadway do
  OpentelemetryTelemetry.end_telemetry_span(@tracer_id, metadata)
  end
 
+ @doc false
+ def handle_batch_start(
+ _event,
+ _measurements,
+ %{
+ topology_name: topology_name,
+ messages: messages,
+ batch_info: %Broadway.BatchInfo{batcher: batcher},
+ } = metadata,
+ _config
+ ) do
+ span_name = "#{inspect(topology_name)}/#{Atom.to_string(batcher)} batching process"
+
+ attributes = %{
+ :"broadway.batch.count" => length(messages)
+ }
+
+ OpentelemetryTelemetry.start_telemetry_span(@tracer_id, span_name, metadata, %{
+ kind: :internal,
+ attributes: attributes
+ })
+ end
+
+ @doc false
+ def handle_batch_stop(
+ _event,
+ _measurements,
+ %{
+ successful_messages: successful_messages,
+ failed_messages: failed_messages
+ } = metadata,
+ _config
+ ) do
+ status = if Enum.empty?(failed_messages) do
+ OpenTelemetry.status(:ok)
+ else
+ OpenTelemetry.status(:error, "Batch completed with failed messages")
+ end
+
+ attributes = %{
+ :"broadway.batch.success.count" => length(successful_messages),
+ :"broadway.batch.failed.count" => length(failed_messages)
+ }
+
+ ctx = OpentelemetryTelemetry.set_current_telemetry_span(@tracer_id, metadata)
+
+ OpenTelemetry.Span.set_status(ctx, status)
+ OpenTelemetry.Span.set_attributes(ctx, attributes)
+
+ OpentelemetryTelemetry.end_telemetry_span(@tracer_id, metadata)
+ end
+
  defp format_error(err) when is_binary(err), do: err
  defp format_error(err), do: inspect(err)
 end

instrumentation/opentelemetry_broadway/test/opentelemetry_broadway_test.exs

@@ -87,4 +87,74 @@ defmodule OpentelemetryBroadwayTest do
  status: ^expected_status
  )}
  end
+
+ test "records span on successful batch" do
+ ref = Broadway.test_batch(TestBroadway, ["batch success", "batch success"])
+
+ assert_receive {:ack, ^ref, [%{data: "batch success"}, %{data: "batch success"}], []}
+
+ expected_status = OpenTelemetry.status(:ok, "")
+
+ assert_receive {:span,
+ span(
+ name: "TestBroadway/default batching process",
+ attributes: attributes,
+ parent_span_id: :undefined,
+ kind: :internal,
+ status: ^expected_status
+ )}
+
+ assert %{
+ "broadway.batch.count": 2,
+ "broadway.batch.failed.count": 0,
+ "broadway.batch.success.count": 2
+ } = :otel_attributes.map(attributes)
+ end
+
+ test "records span on batch with exception" do
+ ref = Broadway.test_batch(TestBroadway, ["batch success", "batch exception"])
+
+ # Uncaught exceptions cause Broadway to fail the whole batch
+ assert_receive {:ack, ^ref, [], [%{data: "batch success"}, %{data: "batch exception"}]}
+
+ expected_status = OpenTelemetry.status(:error, "Batch completed with failed messages")
+
+ assert_receive {:span,
+ span(
+ name: "TestBroadway/default batching process",
+ attributes: attributes,
+ parent_span_id: :undefined,
+ kind: :internal,
+ status: ^expected_status
+ )}
+
+ assert %{
+ "broadway.batch.count": 2,
+ "broadway.batch.failed.count": 2,
+ "broadway.batch.success.count": 0
+ } = :otel_attributes.map(attributes)
+ end
+
+ test "records span on batch with failed messages" do
+ ref = Broadway.test_batch(TestBroadway, ["batch success", "batch error"])
+
+ assert_receive {:ack, ^ref, [%{data: "batch success"}], [%{data: "batch error"}]}
+
+ expected_status = OpenTelemetry.status(:error, "Batch completed with failed messages")
+
+ assert_receive {:span,
+ span(
+ name: "TestBroadway/default batching process",
+ attributes: attributes,
+ parent_span_id: :undefined,
+ kind: :internal,
+ status: ^expected_status
+ )}
+
+ assert %{
+ "broadway.batch.count": 2,
+ "broadway.batch.failed.count": 1,
+ "broadway.batch.success.count": 1
+ } = :otel_attributes.map(attributes)
+ end
 end

instrumentation/opentelemetry_broadway/test/support/test_broadway.ex

@@ -11,22 +11,31 @@ defmodule TestBroadway do
  default: []
  ],
  batchers: [
- default: []
+ # Set a predictable batch_size so our batch tests will work as expected
+ default: [batch_size: 2]
  ]
  )
  end
 
  @impl true
  def handle_message(_processor, %Broadway.Message{} = message, _context) do
  case message.data do
- "success" -> message
  "error" -> Broadway.Message.failed(message, "something went wrong")
  "exception" -> raise RuntimeError, "an exception occurred"
+ _ -> message
  end
  end
 
  @impl true
  def handle_batch(_batcher, messages, _batch_info, _context) do
- messages
+ # Broadway.test_batch/3 will pass `messages` through `handle_message/4` before it hits
+ # handle_batch/3. Thus we need to differentiate a message failed during "batching" vs "intake"
+ Enum.map(messages, fn message ->
+ case message.data do
+ "batch error" -> Broadway.Message.failed(message, "something went wrong")
+ "batch exception" -> raise RuntimeError, "an exception occurred"
+ _ -> message
+ end
+ end)
  end
 end