Push the span context to gracefully handle multiple routers

lib/opentelemetry_plug.ex

@@ -70,9 +70,13 @@ defmodule OpentelemetryPlug do
 
   @doc false
   def handle_start(_, _measurements, %{conn: conn, route: route}, _config) do
-    save_parent_ctx()
-    # setup OpenTelemetry context based on request headers
-    :otel_propagator.text_map_extract(conn.req_headers)
+    parent_ctx = save_parent_ctx()
+
+    if parent_ctx == :undefined do
+      # setup OpenTelemetry context based on request headers, but only if
+      # there's no context already
+      :otel_propagator.text_map_extract(conn.req_headers)
+    end
 
     span_name = "#{route}"
 
@@ -186,12 +190,30 @@ defmodule OpentelemetryPlug do
   @ctx_key {__MODULE__, :parent_ctx}
   defp save_parent_ctx() do
     ctx = Tracer.current_span_ctx()
-    Process.put(@ctx_key, ctx)
+
+    case Process.get(@ctx_key, :undefined) do
+      list when is_list(list) ->
+        Process.put(@ctx_key, [ctx | list])
+
+      :undefined ->
+        Process.put(@ctx_key, [ctx])
+    end
+
+    ctx
   end
 
   defp restore_parent_ctx() do
-    ctx = Process.get(@ctx_key, :undefined)
-    Process.delete(@ctx_key)
+    ctx =
+      case Process.get(@ctx_key, :undefined) do
+        [ctx | rest] ->
+          Process.put(@ctx_key, rest)
+          ctx
+
+        _ ->
+          Process.delete(@ctx_key)
+          :undefined
+      end
+
     Tracer.set_current_span(ctx)
   end
 end

test/opentelemetry_plug_test.exs

@@ -84,6 +84,18 @@ defmodule OpentelemetryPlugTest do
     assert status(code: :error, message: _) = span_status
   end
 
+  test "gracefully handles nested routers" do
+    assert {200, _, _} = request(:get, "/hello/nested")
+
+    assert_receive {:span,
+                    span(name: "/hello/nested/*glob/", parent_span_id: parent, trace_id: trace_id)},
+                   5000
+
+    assert_receive {:span,
+                    span(name: "/hello/nested/*glob", span_id: ^parent, trace_id: ^trace_id)},
+                   5000
+  end
+
   defp base_url do
     info = :ranch.info(MyRouter.HTTP)
     port = Keyword.fetch!(info, :port)
@@ -109,13 +121,26 @@ defmodule OpentelemetryPlugTest do
   end
 end
 
+defmodule MyRouter.NestedRouter do
+  use Plug.Router
+
+  plug :match
+  plug :dispatch
+
+  match "/" do
+    send_resp(conn, 200, "Hello from nested")
+  end
+end
+
 defmodule MyRouter do
   use Plug.Router
 
   plug :match
   plug OpentelemetryPlug.Propagation
   plug :dispatch
 
+  forward "/hello/nested", to: MyRouter.NestedRouter
+
   match "/hello/crash" do
     _ = conn
     raise ArgumentError