From 56c6503a93950353396ede60dd72f4634f881cf4 Mon Sep 17 00:00:00 2001
From: Andrew Rosa <dev@andrewhr.io>
Date: Tue, 15 Nov 2022 21:55:32 -0300
Subject: [PATCH] Fix two use cases of Ecto span propagation (#82)

First one is related to `OpenTelemetry.Ctx` API. I've noticed in a few
scenarios the current span of a trace may get lost after Ecto calls.

Looking at the The `attach/1` typespec, it's a Ctx -> Token, while
`dettach/1` as Token -> Ctx function. That made me assume the expected
input of dettach is the return type of attach. Indeed, after this change
we got the behavior of Ecto calls preserve the parent span untouched.

That leads to a second bug found. When ecto does simple calls within a
Task, due the special propagation code for preloads that means it will
skip the current span, if any.  The solution here is to first check the
current process.

One test was added to reproduce this bug.
---
 .../lib/opentelemetry_ecto.ex                 | 22 +++++++++++++-----
 .../test/opentelemetry_ecto_test.exs          | 23 +++++++++++++++++++
 .../opentelemetry_ecto/test/test_helper.exs   |  2 +-
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/instrumentation/opentelemetry_ecto/lib/opentelemetry_ecto.ex b/instrumentation/opentelemetry_ecto/lib/opentelemetry_ecto.ex
index 44949b7..25d0845 100644
--- a/instrumentation/opentelemetry_ecto/lib/opentelemetry_ecto.ex
+++ b/instrumentation/opentelemetry_ecto/lib/opentelemetry_ecto.ex
@@ -108,11 +108,21 @@ defmodule OpentelemetryEcto do
       |> Map.merge(base_attributes)
       |> Map.merge(additional_attributes)
 
-    parent_context = OpentelemetryProcessPropagator.fetch_parent_ctx(1, :"$callers")
+    parent_context =
+      case OpentelemetryProcessPropagator.fetch_ctx(self()) do
+        :undefined ->
+          OpentelemetryProcessPropagator.fetch_parent_ctx(1, :"$callers")
 
-    if parent_context != :undefined do
-      OpenTelemetry.Ctx.attach(parent_context)
-    end
+        ctx ->
+          ctx
+      end
+
+    parent_token =
+      if parent_context != :undefined do
+        OpenTelemetry.Ctx.attach(parent_context)
+      else
+        :undefined
+      end
 
     s =
       OpenTelemetry.Tracer.start_span(span_name, %{
@@ -131,8 +141,8 @@ defmodule OpentelemetryEcto do
 
     OpenTelemetry.Span.end_span(s)
 
-    if parent_context != :undefined do
-      OpenTelemetry.Ctx.detach(parent_context)
+    if parent_token != :undefined do
+      OpenTelemetry.Ctx.detach(parent_token)
     end
   end
 
diff --git a/instrumentation/opentelemetry_ecto/test/opentelemetry_ecto_test.exs b/instrumentation/opentelemetry_ecto/test/opentelemetry_ecto_test.exs
index 9f6fdec..2c02025 100644
--- a/instrumentation/opentelemetry_ecto/test/opentelemetry_ecto_test.exs
+++ b/instrumentation/opentelemetry_ecto/test/opentelemetry_ecto_test.exs
@@ -192,6 +192,29 @@ defmodule OpentelemetryEctoTest do
     assert_receive {:span, span(parent_span_id: ^root_span_id, name: "opentelemetry_ecto.test_repo.query:comments")}
   end
 
+  test "nested query within Task does not skip parent span" do
+    user = Repo.insert!(%User{email: "opentelemetry@erlang.org"})
+    Repo.insert!(%Post{body: "We got traced!", user: user})
+    Repo.insert!(%Comment{body: "We got traced!", user: user})
+
+    attach_handler()
+
+    Tracer.with_span "root span" do
+      task =
+        Task.async(fn ->
+          Tracer.with_span "parent span" do
+            Repo.all(User)
+          end
+        end)
+
+      Task.await(task)
+    end
+
+    assert_receive {:span, span(span_id: _root_span_id, name: "root span")}
+    assert_receive {:span, span(span_id: parent_span_id, name: "parent span")}
+    assert_receive {:span, span(parent_span_id: ^parent_span_id, name: "opentelemetry_ecto.test_repo.query:users")}
+  end
+
   def attach_handler(config \\ []) do
     # For now setup the handler manually in each test
     handler = {__MODULE__, self()}
diff --git a/instrumentation/opentelemetry_ecto/test/test_helper.exs b/instrumentation/opentelemetry_ecto/test/test_helper.exs
index 6e249dc..7e95f22 100644
--- a/instrumentation/opentelemetry_ecto/test/test_helper.exs
+++ b/instrumentation/opentelemetry_ecto/test/test_helper.exs
@@ -1,5 +1,5 @@
 OpentelemetryEcto.TestRepo.start_link()
 
-ExUnit.start()
+ExUnit.start(capture_log: true)
 
 Ecto.Adapters.SQL.Sandbox.mode(OpentelemetryEcto.TestRepo, {:shared, self()})