Add inject_context and current_context (#78)

* Add inject_context and current_context

* Handle map-based headers in inject_context

* Assume we want to inject the current context in Tracer

* inject_context should pass through the headers if called with no active span

CHANGELOG.md

@@ -5,10 +5,27 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
-- No unreleased changes currently.
+
+[NEXT]: https://github.com/spandex-project/spandex/compare/vNEXT...v2.2.0
+
+### Added
+
+- `Spandex.current_context/1` and `Spandex.Tracer.current_context/1` functions,
+  which get a `Spandex.SpanContext` struct based on the current context.
+
+- `Spandex.inject_context/3` and `Spandex.Tracer.inject_context/2` functions,
+  which inject a distributed tracing context into a list of HTTP headers.
+
+### Changed
+
+- The `Spandex.Adapter` behaviour now requires an `inject_context/3` callback,
+  which encodes a `Spandex.SpanContext` as HTTP headers for distributed
+  tracing.
 
 ## [2.2.0]
 
+[2.2.0]: https://github.com/spandex-project/spandex/compare/v2.2.0...v2.1.0
+
 ### Added
 - The `Spandex.Trace` struct now includes `priority` and `baggage` fields, to
   support priority sampling of distributed traces and trace-level baggage,
@@ -39,6 +56,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [2.1.0]
 It is recommended to reread the README, to see the upgrade guide and understand the changes.
 
+[2.1.0]: https://github.com/spandex-project/spandex/compare/v2.1.0...v1.6.1
+
 ### Added
 - Massive changes, including separating adapters into their own repositories
 
@@ -49,10 +68,16 @@ It is recommended to reread the README, to see the upgrade guide and understand
 - Adapters now exist in their own repositories
 
 ## [1.6.1] - 2018-06-04
+
+[1.6.1]: https://github.com/spandex-project/spandex/compare/v1.6.1...v1.6.0
+
 ### Added
 - `private` key, when updating spans, for non-inheriting meta
 
 ## [1.6.0] - 2018-06-04
+
+[1.6.0]: https://github.com/spandex-project/spandex/compare/v1.6.0...v1.5.0
+
 ### Added
 - Storage strategy behaviour
 

lib/adapter.ex

@@ -6,6 +6,7 @@ defmodule Spandex.Adapter do
   @callback distributed_context(Plug.Conn.t(), Keyword.t()) ::
               {:ok, Spandex.SpanContext.t()}
               | {:error, atom()}
+  @callback inject_context(Spandex.headers(), Spandex.SpanContext.t(), Keyword.t()) :: Spandex.headers()
   @callback trace_id() :: Spandex.id()
   @callback span_id() :: Spandex.id()
   @callback now() :: Spandex.timestamp()

lib/spandex.ex

@@ -11,6 +11,8 @@ defmodule Spandex do
     Tracer
   }
 
+  @type headers :: [{atom, binary}] | [{binary, binary}] | %{binary => binary}
+
   @typedoc "Used for Span and Trace IDs (type defined by adapters)"
   @type id :: term()
 
@@ -241,8 +243,31 @@ defmodule Spandex do
     end
   end
 
+  @spec current_context(Tracer.opts()) ::
+          {:ok, SpanContext.t()}
+          | {:error, :disabled}
+          | {:error, :no_span_context}
+          | {:error, :no_trace_context}
+          | {:error, [Optimal.error()]}
+  def current_context(:disabled), do: {:error, :disabled}
+
+  def current_context(opts) do
+    strategy = opts[:strategy]
+
+    case strategy.get_trace(opts[:trace_key]) do
+      {:ok, %Trace{id: trace_id, priority: priority, baggage: baggage, stack: [%Span{id: span_id} | _]}} ->
+        {:ok, %SpanContext{trace_id: trace_id, priority: priority, baggage: baggage, parent_id: span_id}}
+
+      {:ok, %Trace{stack: []}} ->
+        {:error, :no_span_context}
+
+      {:error, _} ->
+        {:error, :no_trace_context}
+    end
+  end
+
   @spec continue_trace(String.t(), SpanContext.t(), Keyword.t()) ::
-          {:ok, %Trace{}}
+          {:ok, Trace.t()}
           | {:error, :disabled}
           | {:error, :trace_already_present}
           | {:error, [Optimal.error()]}
@@ -260,7 +285,7 @@ defmodule Spandex do
   end
 
   @spec continue_trace(String.t(), Spandex.id(), Spandex.id(), Keyword.t()) ::
-          {:ok, %Trace{}}
+          {:ok, Trace.t()}
           | {:error, :disabled}
           | {:error, :trace_already_present}
           | {:error, [Optimal.error()]}
@@ -272,7 +297,7 @@ defmodule Spandex do
   end
 
   @spec continue_trace_from_span(String.t(), Span.t(), Tracer.opts()) ::
-          {:ok, %Trace{}}
+          {:ok, Trace.t()}
           | {:error, :disabled}
           | {:error, :trace_already_present}
           | {:error, [Optimal.error()]}
@@ -290,8 +315,8 @@ defmodule Spandex do
   end
 
   @spec distributed_context(Plug.Conn.t(), Tracer.opts()) ::
-          {:ok, map()}
-          | {:error, atom()}
+          {:ok, SpanContext.t()}
+          | {:error, :disabled}
           | {:error, [Optimal.error()]}
   def distributed_context(_, :disabled), do: {:error, :disabled}
 
@@ -300,6 +325,12 @@ defmodule Spandex do
     adapter.distributed_context(conn, opts)
   end
 
+  @spec inject_context(headers(), SpanContext.t(), Tracer.opts()) :: headers()
+  def inject_context(headers, %SpanContext{} = span_context, opts) do
+    adapter = opts[:adapter]
+    adapter.inject_context(headers, span_context, opts)
+  end
+
   # Private Helpers
 
   defp do_continue_trace(name, span_context, opts) do

lib/tracer.ex

@@ -34,7 +34,14 @@ defmodule Spandex.Tracer do
   @callback current_trace_id(opts) :: nil | Spandex.id()
   @callback current_span_id(opts) :: nil | Spandex.id()
   @callback current_span(opts) :: nil | Span.t()
+  @callback current_context(opts) ::
+              {:ok, SpanContext.t()}
+              | {:error, :disabled}
+              | {:error, :no_span_context}
+              | {:error, :no_trace_context}
+              | {:error, [Optimal.error()]}
   @callback distributed_context(Plug.Conn.t(), opts) :: tagged_tuple(map)
+  @callback inject_context(Spandex.headers(), opts) :: Spandex.headers()
   @macrocallback span(span_name, opts, do: Macro.t()) :: Macro.t()
   @macrocallback trace(span_name, opts, do: Macro.t()) :: Macro.t()
 
@@ -197,6 +204,7 @@ defmodule Spandex.Tracer do
 
       @impl Spandex.Tracer
       def continue_trace(span_name, span_context, opts \\ [])
+
       def continue_trace(span_name, %SpanContext{} = span_context, opts) do
         Spandex.continue_trace(span_name, span_context, config(opts, @otp_app))
       end
@@ -231,11 +239,28 @@ defmodule Spandex.Tracer do
         Spandex.current_span(config(opts, @otp_app))
       end
 
+      @impl Spandex.Tracer
+      def current_context(opts \\ []) do
+        Spandex.current_context(config(opts, @otp_app))
+      end
+
       @impl Spandex.Tracer
       def distributed_context(conn, opts \\ []) do
         Spandex.distributed_context(conn, config(opts, @otp_app))
       end
 
+      @impl Spandex.Tracer
+      def inject_context(headers, opts \\ []) do
+        opts
+        |> current_context()
+        |> case do
+          {:ok, span_context} ->
+            Spandex.inject_context(headers, span_context, config(opts, @otp_app))
+
+          _ -> headers
+        end
+      end
+
       defp merge_config(opts, otp_app) do
         otp_app
         |> Application.get_env(__MODULE__)

test/spandex_test.exs

@@ -465,14 +465,14 @@ defmodule Spandex.Test.SpandexTest do
       assert span_id == Spandex.current_span_id(@base_opts)
     end
 
-    test "returns nil if no trace is active" do
+    test "returns nil if no span is active" do
       opts = @base_opts ++ @span_opts
       assert {:ok, %Trace{}} = Spandex.start_trace("root_span", opts)
       assert {:ok, %Span{}} = Spandex.finish_span(@base_opts)
       assert nil == Spandex.current_span_id(@base_opts)
     end
 
-    test "returns nil if no span is active" do
+    test "returns nil if no trace is active" do
       assert nil == Spandex.current_span_id(@base_opts)
     end
 
@@ -505,6 +505,30 @@ defmodule Spandex.Test.SpandexTest do
     end
   end
 
+  describe "Spandex.current_context/1" do
+    test "returns the active SpanContext if a span is active" do
+      opts = @base_opts ++ @span_opts
+      assert {:ok, %Trace{id: trace_id}} = Spandex.start_trace("root_span", opts)
+      assert {:ok, %Span{id: span_id}} = Spandex.start_span("span_name", @base_opts)
+      assert {:ok, %SpanContext{trace_id: ^trace_id, parent_id: ^span_id}} = Spandex.current_context(@base_opts)
+    end
+
+    test "returns an error if no span is active" do
+      opts = @base_opts ++ @span_opts
+      assert {:ok, %Trace{}} = Spandex.start_trace("root_span", opts)
+      assert {:ok, %Span{}} = Spandex.finish_span(@base_opts)
+      assert {:error, :no_span_context} == Spandex.current_context(@base_opts)
+    end
+
+    test "returns an error if no trace is active" do
+      assert {:error, :no_trace_context} == Spandex.current_context(@base_opts)
+    end
+
+    test "returns an error if tracing is disabled" do
+      assert {:error, :disabled} == Spandex.current_context(:disabled)
+    end
+  end
+
   describe "Spandex.continue_trace/3" do
     test "starts a new child span in an existing trace based on a specified name, trace ID and parent span ID" do
       opts = @base_opts ++ @span_opts
@@ -622,8 +646,10 @@ defmodule Spandex.Test.SpandexTest do
         |> Plug.Test.conn("/")
         |> Plug.Conn.put_req_header("x-test-trace-id", "1234")
         |> Plug.Conn.put_req_header("x-test-parent-id", "5678")
+        |> Plug.Conn.put_req_header("x-test-sampling-priority", "10")
 
-      assert {:ok, %{trace_id: 1234, parent_id: 5678}} = Spandex.distributed_context(conn, @base_opts)
+      assert {:ok, %SpanContext{} = span_context} = Spandex.distributed_context(conn, @base_opts)
+      assert %SpanContext{trace_id: 1234, parent_id: 5678, priority: 10} = span_context
     end
 
     test "returns an error if distributed tracing headers are not present" do
@@ -636,4 +662,21 @@ defmodule Spandex.Test.SpandexTest do
       assert {:error, :disabled} == Spandex.distributed_context(conn, :disabled)
     end
   end
+
+  describe "Spandex.inject_context/3" do
+    test "Prepends distributed tracing headers to an existing list of headers" do
+      span_context = %SpanContext{trace_id: 123, parent_id: 456, priority: 10}
+      headers = [{"header1", "value1"}, {"header2", "value2"}]
+
+      result = Spandex.inject_context(headers, span_context, @base_opts)
+
+      assert result == [
+               {"x-test-trace-id", "123"},
+               {"x-test-parent-id", "456"},
+               {"x-test-sampling-priority", "10"},
+               {"header1", "value1"},
+               {"header2", "value2"}
+             ]
+    end
+  end
 end

test/support/adapter.ex

@@ -42,6 +42,26 @@ defmodule Spandex.TestAdapter do
     end
   end
 
+  @doc """
+  Injects test HTTP headers to represent the specified SpanContext
+  """
+  @impl Spandex.Adapter
+  @spec inject_context(Spandex.headers(), SpanContext.t(), Tracer.opts()) :: Spandex.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
+
   @spec get_first_header(conn :: Plug.Conn.t(), header_name :: binary) :: binary | nil
   defp get_first_header(conn, header_name) do
     conn
@@ -58,4 +78,12 @@ defmodule Spandex.TestAdapter do
   end
 
   defp parse_header(_header), do: nil
+
+  defp tracing_headers(%SpanContext{trace_id: trace_id, parent_id: parent_id, priority: priority}) do
+    [
+      {"x-test-trace-id", to_string(trace_id)},
+      {"x-test-parent-id", to_string(parent_id)},
+      {"x-test-sampling-priority", to_string(priority)}
+    ]
+  end
 end