Handle map-based headers in inject_context

Author: GregMefford

lib/spandex_datadog/adapter.ex

@@ -53,12 +53,17 @@ defmodule SpandexDatadog.Adapter do
   """
   @impl Spandex.Adapter
   @spec inject_context([{term(), term()}], SpanContext.t(), Tracer.opts()) :: [{term(), term()}]
-  def inject_context(headers, %SpanContext{trace_id: trace_id, parent_id: parent_id, priority: priority}, _opts) do
-    [
-      {"x-datadog-trace-id", to_string(trace_id)},
-      {"x-datadog-parent-id", to_string(parent_id)},
-      {"x-datadog-sampling-priority", to_string(priority)}
-    ] ++ headers
+  def inject_context(headers, %SpanContext{} = span_context, _opts) when is_list(headers) do
+    span_context
+    |> tracing_headers()
+    |> Kernel.++(headers)
+  end
+
+  def inject_context(headers, %SpanContext{} = span_context, _opts) when is_map(headers) do
+    span_context
+    |> tracing_headers()
+    |> Enum.into(%{})
+    |> Map.merge(headers)
   end
 
   # Private Helpers
@@ -79,4 +84,12 @@ defmodule SpandexDatadog.Adapter do
   end
 
   defp parse_header(_header), do: nil
+
+  defp tracing_headers(%SpanContext{trace_id: trace_id, parent_id: parent_id, priority: priority}) do
+    [
+      {"x-datadog-trace-id", to_string(trace_id)},
+      {"x-datadog-parent-id", to_string(parent_id)},
+      {"x-datadog-sampling-priority", to_string(priority)}
+    ]
+  end
 end

test/support/adapter_test.exs

@@ -121,5 +121,20 @@ defmodule Spandex.Test.Datadog.AdapterTest do
                {"header2", "value2"}
              ]
     end
+
+    test "Merges distributed tracing headers with an existing map of headers" do
+      span_context = %SpanContext{trace_id: 123, parent_id: 456, priority: 10}
+      headers = %{"header1" => "value1", "header2" => "value2"}
+
+      result = Adapter.inject_context(headers, span_context, [])
+
+      assert result == %{
+               "x-datadog-trace-id" => "123",
+               "x-datadog-parent-id" => "456",
+               "x-datadog-sampling-priority" => "10",
+               "header1" => "value1",
+               "header2" => "value2"
+             }
+    end
   end
 end