WIP: configurable span tracking behaviour

Author: Garth Kidd

config/config.exs

@@ -8,6 +8,7 @@ if Mix.env() == :test do
     report_dir: "reports/exunit"
 
   config :opencensus,
+    process_context: Opencensus.ProcessContext.DefaultImplementation,
     reporters: [{Opencensus.TestSupport.SpanCaptureReporter, []}],
     send_interval_ms: 100
 end

lib/opencensus/process_context.ex

@@ -0,0 +1,63 @@
+defmodule Opencensus.ProcessContext do
+  @moduledoc """
+  Experimental behaviour to control how the span context is tracked.
+
+  Use the matching methods in `Opencensus.ProcessContext.ConfiguredImplementation`.
+  """
+
+  @doc """
+  Put the current context.
+
+  Returns the previous value.
+  """
+  @callback put_span_ctx(span_ctx :: :opencensus.span_ctx() | :undefined) ::
+              :opencensus.span_ctx() | :undefined
+
+  @doc """
+  Get the current span context.
+
+  Implementations [SHOULD NOT] attempt recovery if the span context isn't where `c:put_span_ctx/1`
+  should have put it. Callers [MAY] rely on this behaviour, using `c:put_span_ctx/1` to "put back"
+  any value they got from `c:get_span_ctx/0`.
+
+  [SHOULD NOT]: https://tools.ietf.org/html/rfc2119#section-4
+  [MAY]: https://tools.ietf.org/html/rfc2119#section-5
+  """
+  @callback get_span_ctx() :: :opencensus.span_ctx() | :undefined
+
+  @doc """
+  Recover the current span context by less reliable means.
+
+  Implementations [SHOULD] check `c:get_span_ctx/0` and return its value if not `:undefined`.
+
+  Callers [SHOULD] check `c:get_span_ctx/0` and avoid calling `c:recover_span_ctx/0` if possible.
+  Callers [MUST NOT] pass a value obtained via `c:recover_span_ctx/0` to `c:put_span_ctx/1`.
+
+  [SHOULD]: https://tools.ietf.org/html/rfc2119#section-3
+  [MUST NOT]: https://tools.ietf.org/html/rfc2119#section-2
+  """
+  @callback recover_span_ctx() :: :opencensus.span_ctx() | :undefined
+end
+
+defmodule Opencensus.ProcessContext.DefaultImplementation do
+  @moduledoc "Process context behaviour matching the default implementation."
+
+  @behaviour Opencensus.ProcessContext
+
+  @impl true
+  def put_span_ctx(span_ctx) do
+    previous_span_ctx = get_span_ctx()
+    Process.put(:oc_span_ctx_key, span_ctx)
+    previous_span_ctx
+  end
+
+  @impl true
+  def get_span_ctx do
+    Process.get(:oc_span_ctx_key, :undefined)
+  end
+
+  @impl true
+  def recover_span_ctx do
+    get_span_ctx()
+  end
+end

lib/opencensus/trace.ex

@@ -45,6 +45,9 @@ defmodule Opencensus.Trace do
   end
   ```
   """
+
+  @behaviour Opencensus.ProcessContext
+
   defmacro with_child_span(label, attributes \\ quote(do: %{}), do: block) do
     line = __CALLER__.line
     module = __CALLER__.module
@@ -60,21 +63,22 @@ defmodule Opencensus.Trace do
       })
 
     quote do
-      parent_span_ctx = :ocp.current_span_ctx()
+      previous_span_ctx = Opencensus.Trace.get_span_ctx()
+      parent_span_ctx = Opencensus.Trace.effective_span_ctx()
 
       new_span_ctx =
         :oc_trace.start_span(unquote(label), parent_span_ctx, %{
           :attributes => unquote(computed_attributes)
         })
 
-      _ = :ocp.with_span_ctx(new_span_ctx)
+      _ = Opencensus.Trace.put_span_ctx(new_span_ctx)
       Opencensus.Logger.set_logger_metadata()
 
       try do
         unquote(block)
       after
         _ = :oc_trace.finish_span(new_span_ctx)
-        _ = :ocp.with_span_ctx(parent_span_ctx)
+        _ = Opencensus.Trace.put_span_ctx(previous_span_ctx)
         Opencensus.Logger.set_logger_metadata()
       end
     end
@@ -166,4 +170,62 @@ defmodule Opencensus.Trace do
   """
   @spec await(Task.t(), :infinity | pos_integer()) :: term()
   defdelegate await(task, timeout \\ 5000), to: Task
+
+  @doc """
+  Put the current span context.
+
+  Replaces `:ocp.with_span_ctx/1`.
+
+  Callers [MAY] pass values from `get_span_ctx/0` to `put_span_ctx/1`.
+
+  Callers [MUST NOT] pass a value obtained via `recover_span_ctx/0` to `put_span_ctx/1`.
+
+  Uses the configured `process_context`. See also: `Opencensus.ProcessContext`.
+
+  [MAY]: https://tools.ietf.org/html/rfc2119#section-5
+  """
+  @impl Opencensus.ProcessContext
+  def put_span_ctx(span_ctx), do: process_context() |> apply(:put_span_ctx, [span_ctx])
+
+  @doc """
+  Get the current span context.
+
+  Replaces `:ocp.current_span_ctx/0`, along with `recover_span_ctx/0`.
+
+  Callers [MAY] pass values from `get_span_ctx/0` to `put_span_ctx/1`.
+
+  Uses the configured `process_context`. See also: `Opencensus.ProcessContext`.
+
+  [MAY]: https://tools.ietf.org/html/rfc2119#section-5
+  """
+  @impl Opencensus.ProcessContext
+  def get_span_ctx, do: process_context() |> apply(:get_span_ctx, [])
+
+  @doc """
+  Recover the current span context by less reliable means.
+
+  Replaces `:ocp.current_span_ctx/0`, along with `get_span_ctx/0`.
+
+  Callers [MUST NOT] pass a value obtained via `recover_span_ctx/0` to `put_span_ctx/1`.
+
+  Uses the configured `process_context`. See also: `Opencensus.ProcessContext`.
+  """
+  @impl Opencensus.ProcessContext
+  def recover_span_ctx, do: process_context() |> apply(:recover_span_ctx, [])
+
+  @doc false
+  def effective_span_ctx do
+    case get_span_ctx() do
+      :undefined -> recover_span_ctx()
+      span_ctx -> span_ctx
+    end
+  end
+
+  defp process_context do
+    Application.get_env(
+      :opencensus,
+      :process_context,
+      Opencensus.ProcessContext.DefaultImplementation
+    )
+  end
 end

test/opencensus_trace_async_test.exs

@@ -42,7 +42,7 @@ defmodule Opencensus.AsyncTest do
     {inner, outer} =
       Trace.with_child_span "outside" do
         outer = :ocp.current_span_ctx() |> Span.load()
-        Trace.async(M, :f, [outer]) |> Trace.await(10)
+        M |> Trace.async(:f, [outer]) |> Trace.await(10)
       end
 
     assert inner.trace_id == outer.trace_id