Distributed Tracing (#36)

Continue trace when trace headers are present

lib/adapters/adapter.ex

@@ -17,4 +17,5 @@ defmodule Spandex.Adapters.Adapter do
   @callback continue_trace_from_span(String.t, map) :: {:ok, term} | {:error, term}
   @callback update_top_span(map) :: :ok | {:error, term}
   @callback update_all_spans(map) :: :ok | {}
+  @callback distributed_context(Plug.Conn.t) :: {:ok, term} | {:error, term}
 end

lib/adapters/datadog.ex

@@ -5,6 +5,7 @@ defmodule Spandex.Adapters.Datadog do
 
   @behaviour Spandex.Adapters.Adapter
 
+  alias Spandex.Adapters.Helpers
   alias Spandex.Datadog.Api
   alias Spandex.Datadog.Span
   alias Spandex.Datadog.Utils
@@ -18,6 +19,7 @@ defmodule Spandex.Adapters.Datadog do
   def start_trace(name) do
     if get_trace() do
       Logger.error("Tried to start a trace over top of another trace.")
+      {:error, :trace_running}
     else
       trace_id = Utils.next_id()
       top_span =
@@ -253,6 +255,22 @@ defmodule Spandex.Adapters.Datadog do
     update_span(%{error: 1, error_message: message, stacktrace: stacktrace, error_type: type})
   end
 
+  @doc """
+  Fetches the datadog trace & parent IDs from the conn request headers
+  if they are present.
+  """
+  @spec distributed_context(conn :: Plug.Conn.t) :: {:ok, %{trace_id: binary, parent_id: binary}} | {:error, :no_trace_context}
+  def distributed_context(%Plug.Conn{} = conn) do
+    trace_id = Helpers.get_first_header(conn, "x-datadog-trace-id")
+    parent_id = Helpers.get_first_header(conn, "x-datadog-parent-id")
+
+    if is_nil(trace_id) || is_nil(parent_id) do
+      {:error, :no_distributed_trace}
+    else
+      {:ok, %{trace_id: trace_id, parent_id: parent_id}}
+    end
+  end
+
   @spec get_trace(term) :: term
   defp get_trace(default \\ nil) do
     Process.get(:spandex_trace, default)

lib/adapters/helpers.ex

@@ -35,4 +35,11 @@ defmodule Spandex.Adapters.Helpers do
 
     configured_level_position <= span_level_position
   end
+
+  @spec get_first_header(conn :: Plug.Conn.t, header_name :: binary) :: binary | nil
+  def get_first_header(conn, header_name) do
+    conn
+    |> Plug.Conn.get_req_header(header_name)
+    |> List.first()
+  end
 end

lib/plug/start_trace.ex

@@ -15,11 +15,23 @@ defmodule Spandex.Plug.StartTrace do
     if ignoring_request?(conn) do
       Utils.trace(conn, false)
     else
-      Spandex.start_trace("request", %{level: Spandex.highest_level()})
-      Utils.trace(conn, true)
+      begin_tracing(conn)
     end
   end
 
+  @spec begin_tracing(conn :: Plug.Conn.t) :: Plug.Conn.t
+  defp begin_tracing(conn) do
+    case Spandex.distributed_context(conn) do
+      {:ok, %{trace_id: trace_id, parent_id: parent_id}} ->
+        Spandex.continue_trace("request", trace_id, parent_id)
+
+      {:error, :no_distributed_trace} ->
+        Spandex.start_trace("request", %{level: Spandex.highest_level()})
+    end
+
+    Utils.trace(conn, true)
+  end
+
   @spec ignoring_request?(conn :: Plug.Conn.t) :: boolean
   defp ignoring_request?(conn) do
     disabled?() || ignored_method?(conn) || ignored_route?(conn)

lib/spandex.ex

@@ -98,4 +98,5 @@ defmodule Spandex do
   delegate_to_adapter(:current_trace_id, [])
   delegate_to_adapter(:current_span_id, [])
   delegate_to_adapter(:current_span, [])
+  delegate_to_adapter(:distributed_context, [conn])
 end

test/plug/start_trace_test.exs

@@ -71,6 +71,22 @@ defmodule Spandex.Plug.StartTraceTest do
       end
     end
 
+    test "continues existing trace when distributed context exists", %{conn: conn} do
+      conn =
+        conn
+        |> Plug.Conn.put_req_header("x-datadog-trace-id", "12345")
+        |> Plug.Conn.put_req_header("x-datadog-parent-id", "67890")
+
+      new_conn = StartTrace.call(conn, [])
+
+      assert %{trace_id: "12345", parent_id: "67890"} = Spandex.current_span()
+
+      refute Spandex.current_span_id() == "67890"
+      refute is_nil(Spandex.current_span_id())
+
+      assert new_conn.assigns[:spandex_trace_request?]
+    end
+
     test "starts new trace", %{conn: conn} do
       new_conn = StartTrace.call(conn, [])