Add root_span implementation

sl0thentr0py · sl0thentr0py · commit 56ef07925f10 · 2024-09-10T14:01:06.000+02:00
* add and track underlying root span in a hidden attr `_sentry_root_otel_span` in all subsequent children spans in the span processor on start
* wrap this underlying root otel span in a `POTelSpan` to act as a proxy
* make `POTelSpan` constructor work with explicitly passed in `otel_span`
* implement `__eq__` on `POTelSpan` to make sure proxies to the same
  underlying `_otel_span` are considered the same
diff --git a/sentry_sdk/integrations/opentelemetry/potel_span_processor.py b/sentry_sdk/integrations/opentelemetry/potel_span_processor.py
@@ -1,7 +1,13 @@
 from collections import deque, defaultdict
 from typing import cast
 
-from opentelemetry.trace import format_trace_id, format_span_id
+from opentelemetry.trace import (
+    format_trace_id,
+    format_span_id,
+    get_current_span,
+    INVALID_SPAN,
+    Span as TraceApiSpan,
+)
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import Span, ReadableSpan, SpanProcessor
 
@@ -44,7 +50,8 @@ def __init__(self):
 
     def on_start(self, span, parent_context=None):
         # type: (Span, Optional[Context]) -> None
-        pass
+        if not is_sentry_span(span):
+            self._add_root_span(span, get_current_span(parent_context))
 
     def on_end(self, span):
         # type: (ReadableSpan) -> None
@@ -68,6 +75,21 @@ def force_flush(self, timeout_millis=30000):
         # type: (int) -> bool
         return True
 
+    def _add_root_span(self, span, parent_span):
+        # type: (Span, TraceApiSpan) -> None
+        """
+        This is required to make POTelSpan.root_span work
+        since we can't traverse back to the root purely with otel efficiently.
+        """
+        if parent_span != INVALID_SPAN and not parent_span.get_span_context().is_remote:
+            # child span points to parent's root or parent
+            span._sentry_root_otel_span = getattr(
+                parent_span, "_sentry_root_otel_span", parent_span
+            )
+        else:
+            # root span points to itself
+            span._sentry_root_otel_span = span
+
     def _flush_root_span(self, span):
         # type: (ReadableSpan) -> None
         transaction_event = self._root_span_to_transaction_event(span)
diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -3,7 +3,7 @@
 from datetime import datetime, timedelta, timezone
 
 from opentelemetry import trace as otel_trace, context
-from opentelemetry.trace import format_trace_id, format_span_id
+from opentelemetry.trace import format_trace_id, format_span_id, Span as OtelSpan
 from opentelemetry.trace.status import StatusCode
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -17,7 +17,7 @@
     nanosecond_time,
 )
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping, MutableMapping
@@ -1201,32 +1201,42 @@ def __init__(
         start_timestamp=None,  # type: Optional[Union[datetime, float]]
         origin=None,  # type: Optional[str]
         name=None,  # type: Optional[str]
+        otel_span=None,  # type: Optional[OtelSpan]
         **_,  # type: dict[str, object]
     ):
         # type: (...) -> None
         """
         For backwards compatibility with old the old Span interface, this class
         accepts arbitrary keyword arguments, in addition to the ones explicitly
         listed in the signature. These additional arguments are ignored.
+
+        If otel_span is passed explicitly, just acts as a proxy.
         """
-        from sentry_sdk.integrations.opentelemetry.utils import (
-            convert_to_otel_timestamp,
-        )
+        if otel_span is not None:
+            self._otel_span = otel_span
+        else:
+            from sentry_sdk.integrations.opentelemetry.utils import (
+                convert_to_otel_timestamp,
+            )
 
-        if start_timestamp is not None:
-            # OTel timestamps have nanosecond precision
-            start_timestamp = convert_to_otel_timestamp(start_timestamp)
+            if start_timestamp is not None:
+                # OTel timestamps have nanosecond precision
+                start_timestamp = convert_to_otel_timestamp(start_timestamp)
 
-        self._otel_span = tracer.start_span(
-            description or op or "", start_time=start_timestamp
-        )
+            self._otel_span = tracer.start_span(
+                description or op or "", start_time=start_timestamp
+            )
 
-        self.origin = origin or DEFAULT_SPAN_ORIGIN
-        self.op = op
-        self.description = description
-        self.name = name
-        if status is not None:
-            self.set_status(status)
+            self.origin = origin or DEFAULT_SPAN_ORIGIN
+            self.op = op
+            self.description = description
+            self.name = name
+            if status is not None:
+                self.set_status(status)
+
+    def __eq__(self, other):
+        # type: (POTelSpan) -> bool
+        return self._otel_span == other._otel_span
 
     def __repr__(self):
         # type: () -> str
@@ -1318,17 +1328,16 @@ def containing_transaction(self):
     @property
     def root_span(self):
         # type: () -> Optional[POTelSpan]
-        # XXX implement this
-        # there's a span.parent property, but it returns the parent spancontext
-        # not sure if there's a way to retrieve the parent with pure otel.
-        return None
+        root_otel_span = cast(
+            "Optional[OtelSpan]",
+            getattr(self._otel_span, "_sentry_root_otel_span", None),
+        )
+        return POTelSpan(otel_span=root_otel_span) if root_otel_span else None
 
     @property
     def is_root_span(self):
         # type: () -> bool
-        return (
-            isinstance(self._otel_span, ReadableSpan) and self._otel_span.parent is None
-        )
+        return self.root_span == self
 
     @property
     def parent_span_id(self):
diff --git a/tests/integrations/opentelemetry/test_potel.py b/tests/integrations/opentelemetry/test_potel.py
@@ -314,3 +314,18 @@ def test_multiple_transaction_tags_isolation_scope_started_with_sentry(
 
     assert payload_a["tags"] == {"tag.global": 99, "tag.inner.a": "a"}
     assert payload_b["tags"] == {"tag.global": 99, "tag.inner.b": "b"}
+
+
+def test_potel_span_root_span_references():
+    with sentry_sdk.start_span(description="request") as request_span:
+        assert request_span.is_root_span
+        assert request_span.root_span == request_span
+        with sentry_sdk.start_span(description="db") as db_span:
+            assert not db_span.is_root_span
+            assert db_span.root_span == request_span
+            with sentry_sdk.start_span(description="redis") as redis_span:
+                assert not redis_span.is_root_span
+                assert redis_span.root_span == request_span
+        with sentry_sdk.start_span(description="http") as http_span:
+            assert not http_span.is_root_span
+            assert http_span.root_span == request_span
