New Scope implementation based on OTel Context

sentry_sdk/api.py

@@ -4,9 +4,10 @@
 
 from sentry_sdk import tracing, tracing_utils, Client
 from sentry_sdk._init_implementation import init
-from sentry_sdk.scope import Scope, _ScopeManager, new_scope, isolation_scope
 from sentry_sdk.tracing import POTelSpan, Transaction, trace
 from sentry_sdk.crons import monitor
+# TODO-neel-potel make 2 scope strategies/impls and switch
+from sentry_sdk.integrations.opentelemetry.scope import PotelScope as Scope, new_scope, isolation_scope
 
 from sentry_sdk._types import TYPE_CHECKING
 

sentry_sdk/integrations/opentelemetry/__init__.py

@@ -1,7 +1,8 @@
-from sentry_sdk.integrations.opentelemetry.span_processor import (  # noqa: F401
-    SentrySpanProcessor,
-)
+# TODO-neel-potel fix circular imports
+# from sentry_sdk.integrations.opentelemetry.span_processor import (  # noqa: F401
+#     SentrySpanProcessor,
+# )
 
-from sentry_sdk.integrations.opentelemetry.propagator import (  # noqa: F401
-    SentryPropagator,
-)
+# from sentry_sdk.integrations.opentelemetry.propagator import (  # noqa: F401
+#     SentryPropagator,
+# )

sentry_sdk/integrations/opentelemetry/consts.py

@@ -1,8 +1,14 @@
 from opentelemetry.context import create_key
 
 
+# propagation keys
 SENTRY_TRACE_KEY = create_key("sentry-trace")
 SENTRY_BAGGAGE_KEY = create_key("sentry-baggage")
+
+# scope management keys
+SENTRY_SCOPES_KEY = create_key("sentry_scopes")
+SENTRY_FORK_ISOLATION_SCOPE_KEY = create_key("sentry_fork_isolation_scope")
+
 OTEL_SENTRY_CONTEXT = "otel"
 SPAN_ORIGIN = "auto.otel"
 

sentry_sdk/integrations/opentelemetry/contextvars_context.py

@@ -1,26 +1,32 @@
-from opentelemetry.context import Context, create_key, get_value, set_value
+from opentelemetry.context import Context, get_value, set_value
 from opentelemetry.context.contextvars_context import ContextVarsRuntimeContext
 
-from sentry_sdk.scope import Scope
-
-
-_SCOPES_KEY = create_key("sentry_scopes")
+import sentry_sdk
+from sentry_sdk.integrations.opentelemetry.consts import (
+    SENTRY_SCOPES_KEY,
+    SENTRY_FORK_ISOLATION_SCOPE_KEY,
+)
 
 
 class SentryContextVarsRuntimeContext(ContextVarsRuntimeContext):
     def attach(self, context):
         # type: (Context) -> object
-        scopes = get_value(_SCOPES_KEY, context)
+        scopes = get_value(SENTRY_SCOPES_KEY, context)
+        should_fork_isolation_scope = context.pop(
+            SENTRY_FORK_ISOLATION_SCOPE_KEY, False
+        )
 
         if scopes and isinstance(scopes, tuple):
             (current_scope, isolation_scope) = scopes
         else:
-            current_scope = Scope.get_current_scope()
-            isolation_scope = Scope.get_isolation_scope()
+            current_scope = sentry_sdk.get_current_scope()
+            isolation_scope = sentry_sdk.get_isolation_scope()
 
-        # TODO-neel-potel fork isolation_scope too like JS
-        # once we setup our own apis to pass through to otel
-        new_scopes = (current_scope.fork(), isolation_scope)
-        new_context = set_value(_SCOPES_KEY, new_scopes, context)
+        new_scope = current_scope.fork()
+        new_isolation_scope = (
+            isolation_scope.fork() if should_fork_isolation_scope else isolation_scope
+        )
+        new_scopes = (new_scope, new_isolation_scope)
 
+        new_context = set_value(SENTRY_SCOPES_KEY, new_scopes, context)
         return super().attach(new_context)

sentry_sdk/integrations/opentelemetry/integration.py

@@ -6,7 +6,12 @@
 
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.integrations.opentelemetry.propagator import SentryPropagator
-from sentry_sdk.integrations.opentelemetry.span_processor import SentrySpanProcessor
+from sentry_sdk.integrations.opentelemetry.potel_span_processor import (
+    PotelSentrySpanProcessor,
+)
+from sentry_sdk.integrations.opentelemetry.contextvars_context import (
+    SentryContextVarsRuntimeContext,
+)
 from sentry_sdk.utils import logger
 
 try:
@@ -46,9 +51,14 @@ def setup_once():
 
 def _setup_sentry_tracing():
     # type: () -> None
+    import opentelemetry.context
+
+    opentelemetry.context._RUNTIME_CONTEXT = SentryContextVarsRuntimeContext()
+
     provider = TracerProvider()
-    provider.add_span_processor(SentrySpanProcessor())
+    provider.add_span_processor(PotelSentrySpanProcessor())
     trace.set_tracer_provider(provider)
+
     set_global_textmap(SentryPropagator())
 
 

sentry_sdk/integrations/opentelemetry/scope.py

@@ -0,0 +1,84 @@
+from typing import cast
+from contextlib import contextmanager
+
+from opentelemetry.context import get_value, set_value, attach, detach, get_current
+
+from sentry_sdk.scope import Scope, ScopeType
+from sentry_sdk.integrations.opentelemetry.consts import (
+    SENTRY_SCOPES_KEY,
+    SENTRY_FORK_ISOLATION_SCOPE_KEY,
+)
+
+from sentry_sdk._types import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Tuple, Optional, Generator
+
+
+class PotelScope(Scope):
+    @classmethod
+    def _get_scopes(cls):
+        # type: () -> Optional[Tuple[Scope, Scope]]
+        """
+        Returns the current scopes tuple on the otel context. Internal use only.
+        """
+        return cast("Optional[Tuple[Scope, Scope]]", get_value(SENTRY_SCOPES_KEY))
+
+    @classmethod
+    def get_current_scope(cls):
+        # type: () -> Scope
+        """
+        Returns the current scope.
+        """
+        return cls._get_current_scope() or _INITIAL_CURRENT_SCOPE
+
+    @classmethod
+    def _get_current_scope(cls):
+        # type: () -> Optional[Scope]
+        """
+        Returns the current scope without creating a new one. Internal use only.
+        """
+        scopes = cls._get_scopes()
+        return scopes[0] if scopes else None
+
+    @classmethod
+    def get_isolation_scope(cls):
+        """
+        Returns the isolation scope.
+        """
+        # type: () -> Scope
+        return cls._get_isolation_scope() or _INITIAL_ISOLATION_SCOPE
+
+    @classmethod
+    def _get_isolation_scope(cls):
+        # type: () -> Optional[Scope]
+        """
+        Returns the isolation scope without creating a new one. Internal use only.
+        """
+        scopes = cls._get_scopes()
+        return scopes[1] if scopes else None
+
+
+_INITIAL_CURRENT_SCOPE = PotelScope(ty=ScopeType.CURRENT)
+_INITIAL_ISOLATION_SCOPE = PotelScope(ty=ScopeType.ISOLATION)
+
+
+@contextmanager
+def isolation_scope():
+    # type: () -> Generator[Scope, None, None]
+    context = set_value(SENTRY_FORK_ISOLATION_SCOPE_KEY, True)
+    token = attach(context)
+    try:
+        yield PotelScope.get_isolation_scope()
+    finally:
+        detach(token)
+
+
+@contextmanager
+def new_scope():
+    # type: () -> Generator[Scope, None, None]
+    token = attach(get_current())
+    try:
+        yield PotelScope.get_current_scope()
+    finally:
+        detach(token)

sentry_sdk/scope.py

@@ -255,13 +255,21 @@ def get_current_scope(cls):
 
         Returns the current scope.
         """
-        current_scope = _current_scope.get()
+        current_scope = cls._get_current_scope()
         if current_scope is None:
             current_scope = Scope(ty=ScopeType.CURRENT)
             _current_scope.set(current_scope)
 
         return current_scope
 
+    @classmethod
+    def _get_current_scope(cls):
+        # type: () -> Optional[Scope]
+        """
+        Returns the current scope without creating a new one. Internal use only.
+        """
+        return _current_scope.get()
+
     @classmethod
     def set_current_scope(cls, new_current_scope):
         # type: (Scope) -> None
@@ -281,13 +289,21 @@ def get_isolation_scope(cls):
 
         Returns the isolation scope.
         """
-        isolation_scope = _isolation_scope.get()
+        isolation_scope = cls._get_isolation_scope()
         if isolation_scope is None:
             isolation_scope = Scope(ty=ScopeType.ISOLATION)
             _isolation_scope.set(isolation_scope)
 
         return isolation_scope
 
+    @classmethod
+    def _get_isolation_scope(cls):
+        # type: () -> Optional[Scope]
+        """
+        Returns the isolation scope without creating a new one. Internal use only.
+        """
+        return _isolation_scope.get()
+
     @classmethod
     def set_isolation_scope(cls, new_isolation_scope):
         # type: (Scope) -> None
@@ -342,13 +358,11 @@ def _merge_scopes(self, additional_scope=None, additional_scope_kwargs=None):
         final_scope = copy(_global_scope) if _global_scope is not None else Scope()
         final_scope._type = ScopeType.MERGED
 
-        isolation_scope = _isolation_scope.get()
-        if isolation_scope is not None:
-            final_scope.update_from_scope(isolation_scope)
+        isolation_scope = self.get_isolation_scope()
+        final_scope.update_from_scope(isolation_scope)
 
-        current_scope = _current_scope.get()
-        if current_scope is not None:
-            final_scope.update_from_scope(current_scope)
+        current_scope = self.get_current_scope()
+        final_scope.update_from_scope(current_scope)
 
         if self != current_scope and self != isolation_scope:
             final_scope.update_from_scope(self)
@@ -374,7 +388,7 @@ def get_client(cls):
         This checks the current scope, the isolation scope and the global scope for a client.
         If no client is available a :py:class:`sentry_sdk.client.NonRecordingClient` is returned.
         """
-        current_scope = _current_scope.get()
+        current_scope = cls._get_current_scope()
         try:
             client = current_scope.client
         except AttributeError:
@@ -383,7 +397,7 @@ def get_client(cls):
         if client is not None and client.is_active():
             return client
 
-        isolation_scope = _isolation_scope.get()
+        isolation_scope = cls._get_isolation_scope()
         try:
             client = isolation_scope.client
         except AttributeError:
@@ -1361,8 +1375,8 @@ def run_event_processors(self, event, hint):
 
         if not is_check_in:
             # Get scopes without creating them to prevent infinite recursion
-            isolation_scope = _isolation_scope.get()
-            current_scope = _current_scope.get()
+            isolation_scope = self._get_isolation_scope()
+            current_scope = self._get_current_scope()
 
             event_processors = chain(
                 global_event_processors,

tests/conftest.py

@@ -63,6 +63,7 @@ def benchmark():
 
 
 from sentry_sdk import scope
+import sentry_sdk.integrations.opentelemetry.scope as potel_scope
 
 
 @pytest.fixture(autouse=True)
@@ -74,6 +75,9 @@ def clean_scopes():
     scope._isolation_scope.set(None)
     scope._current_scope.set(None)
 
+    potel_scope._INITIAL_CURRENT_SCOPE.clear()
+    potel_scope._INITIAL_ISOLATION_SCOPE.clear()
+
 
 @pytest.fixture(autouse=True)
 def internal_exceptions(request):

tests/integrations/opentelemetry/test_potel.py

@@ -250,3 +250,67 @@ def test_span_data_started_with_sentry(capture_envelopes):
         "sentry.description": "statement",
         "sentry.op": "db",
     }
+
+
+def test_transaction_tags_started_with_otel(capture_envelopes):
+    envelopes = capture_envelopes()
+
+    sentry_sdk.set_tag("tag.global", 99)
+    with tracer.start_as_current_span("request"):
+        sentry_sdk.set_tag("tag.inner", "foo")
+
+    (envelope,) = envelopes
+    (item,) = envelope.items
+    payload = item.payload.json
+
+    assert payload["tags"] == {"tag.global": 99, "tag.inner": "foo"}
+
+
+def test_transaction_tags_started_with_sentry(capture_envelopes):
+    envelopes = capture_envelopes()
+
+    sentry_sdk.set_tag("tag.global", 99)
+    with sentry_sdk.start_span(description="request"):
+        sentry_sdk.set_tag("tag.inner", "foo")
+
+    (envelope,) = envelopes
+    (item,) = envelope.items
+    payload = item.payload.json
+
+    assert payload["tags"] == {"tag.global": 99, "tag.inner": "foo"}
+
+
+def test_multiple_transaction_tags_isolation_scope_started_with_otel(capture_envelopes):
+    envelopes = capture_envelopes()
+
+    sentry_sdk.set_tag("tag.global", 99)
+    with sentry_sdk.isolation_scope():
+        with tracer.start_as_current_span("request a"):
+            sentry_sdk.set_tag("tag.inner.a", "a")
+    with sentry_sdk.isolation_scope():
+        with tracer.start_as_current_span("request b"):
+            sentry_sdk.set_tag("tag.inner.b", "b")
+
+    (payload_a, payload_b) = [envelope.items[0].payload.json for envelope in envelopes]
+
+    assert payload_a["tags"] == {"tag.global": 99, "tag.inner.a": "a"}
+    assert payload_b["tags"] == {"tag.global": 99, "tag.inner.b": "b"}
+
+
+def test_multiple_transaction_tags_isolation_scope_started_with_sentry(
+    capture_envelopes,
+):
+    envelopes = capture_envelopes()
+
+    sentry_sdk.set_tag("tag.global", 99)
+    with sentry_sdk.isolation_scope():
+        with sentry_sdk.start_span(description="request a"):
+            sentry_sdk.set_tag("tag.inner.a", "a")
+    with sentry_sdk.isolation_scope():
+        with sentry_sdk.start_span(description="request b"):
+            sentry_sdk.set_tag("tag.inner.b", "b")
+
+    (payload_a, payload_b) = [envelope.items[0].payload.json for envelope in envelopes]
+
+    assert payload_a["tags"] == {"tag.global": 99, "tag.inner.a": "a"}
+    assert payload_b["tags"] == {"tag.global": 99, "tag.inner.b": "b"}