Add generic handlers of distributed trace headers

lib/new_relic/distributed_trace.ex

@@ -9,13 +9,20 @@ defmodule NewRelic.DistributedTrace do
   alias NewRelic.Harvest.Collector.AgentRun
   alias NewRelic.Transaction
 
-  def start(:http, headers) do
-    determine_context(:http, headers)
-    |> track_transaction(transport_type: "HTTP")
+  defguardp is_known_transport(transport) when transport in ~w(http generic)a
+
+  def start(transport, headers_or_payload) when is_known_transport(transport) do
+    transport_type =
+      if transport == :http,
+        do: "HTTP",
+        else: Map.get(headers_or_payload, :transport_type, "Unknown")
+
+    determine_context(transport, headers_or_payload)
+    |> track_transaction(transport_type: transport_type)
   end
 
-  defp determine_context(:http, headers) do
-    case accept_distributed_trace_headers(:http, headers) do
+  defp determine_context(transport, headers) when is_known_transport(transport) do
+    case accept_distributed_trace_headers(transport, headers) do
       %Context{} = context -> context
       _ -> generate_new_context()
     end
@@ -25,6 +32,10 @@ defmodule NewRelic.DistributedTrace do
     w3c_headers(headers) || newrelic_header(headers) || :no_payload
   end
 
+  def accept_distributed_trace_headers(:generic, payload) do
+    w3c_headers(payload) || :no_payload
+  end
+
   defp w3c_headers(headers) do
     with {:ok, _traceparent} <- Map.fetch(headers, @w3c_traceparent),
          %Context{} = context <- __MODULE__.W3CTraceContext.extract(headers) do
@@ -43,26 +54,35 @@ defmodule NewRelic.DistributedTrace do
     end
   end
 
-  def distributed_trace_headers(:http) do
+  def distributed_trace_headers(transport) do
     case get_tracing_context() do
       nil ->
         []
 
       context ->
-        context = %{
-          context
-          | span_guid: get_current_span_guid(),
-            timestamp: System.system_time(:millisecond)
-        }
-
-        nr_header = NewRelic.DistributedTrace.NewRelicContext.generate(context)
-        {traceparent, tracestate} = NewRelic.DistributedTrace.W3CTraceContext.generate(context)
-
-        [
-          {@nr_header, nr_header},
-          {@w3c_traceparent, traceparent},
-          {@w3c_tracestate, tracestate}
-        ]
+        context_to_distributed_trace_headers(context, transport)
+    end
+  end
+
+  defp context_to_distributed_trace_headers(context, transport) do
+    context = %{
+      context
+      | span_guid: get_current_span_guid(),
+        timestamp: System.system_time(:millisecond)
+    }
+
+    {traceparent, tracestate} = NewRelic.DistributedTrace.W3CTraceContext.generate(context)
+
+    generic_headers = [
+      {@w3c_traceparent, traceparent},
+      {@w3c_tracestate, tracestate}
+    ]
+
+    if transport == :generic do
+      generic_headers
+    else
+      nr_header = NewRelic.DistributedTrace.NewRelicContext.generate(context)
+      [{@nr_header, nr_header} | generic_headers]
     end
   end
 

test/distributed_trace_test.exs

@@ -66,7 +66,7 @@ defmodule DistributedTraceTest do
     :ok
   end
 
-  test "Annotate Transaction event with DT attrs" do
+  test "Annotate HTTP transaction event with DT attrs" do
     TestHelper.restart_harvest_cycle(Collector.TransactionEvent.HarvestCycle)
 
     TestHelper.request(
@@ -188,6 +188,28 @@ defmodule DistributedTraceTest do
     |> Task.await()
   end
 
+  test "Attaches generic DT attributes to the context" do
+    tracestate =
+      "2534120@nr=0-0-2534120-1651653516-34cc430c00000000-268a50db3d00fd03-1-1.049247-1695854781262"
+
+    traceid = "b9d7bbdbaf2b3725f4b4799f8dd60081"
+    traceparent = "00-#{traceid}-34cc430c00000000-01"
+
+    Task.async(fn ->
+      NewRelic.start_transaction("TransactionCategory", "MyTaskName")
+
+      NewRelic.DistributedTrace.start(:generic, %{
+        "tracestate" => tracestate,
+        "traceparent" => traceparent
+      })
+
+      context = NewRelic.DistributedTrace.get_tracing_context()
+
+      assert context.trace_id == traceid
+    end)
+    |> Task.await()
+  end
+
   describe "Context decoding" do
     test "ignore unknown version" do
       payload = %{"v" => [666]} |> Jason.encode!() |> Base.encode64()