Fix rust tracing

Author: sl0thentr0py

sentry_sdk/integrations/rust_tracing.py

@@ -37,11 +37,9 @@
 import sentry_sdk
 from sentry_sdk.integrations import Integration
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.tracing import Span as SentrySpan
+from sentry_sdk.tracing import POTelSpan as SentrySpan
 from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE
 
-TraceState = Optional[Tuple[Optional[SentrySpan], SentrySpan]]
-
 
 class RustTracingLevel(Enum):
     Trace: str = "TRACE"
@@ -171,7 +169,7 @@ def _include_tracing_fields(self) -> bool:
             else self.include_tracing_fields
         )
 
-    def on_event(self, event: str, _span_state: TraceState) -> None:
+    def on_event(self, event: str, _span_state: Optional[SentrySpan]) -> None:
         deserialized_event = json.loads(event)
         metadata = deserialized_event.get("metadata", {})
 
@@ -185,7 +183,7 @@ def on_event(self, event: str, _span_state: TraceState) -> None:
         elif event_type == EventTypeMapping.Event:
             process_event(deserialized_event)
 
-    def on_new_span(self, attrs: str, span_id: str) -> TraceState:
+    def on_new_span(self, attrs: str, span_id: str) -> Optional[SentrySpan]:
         attrs = json.loads(attrs)
         metadata = attrs.get("metadata", {})
 
@@ -205,48 +203,31 @@ def on_new_span(self, attrs: str, span_id: str) -> TraceState:
         else:
             sentry_span_name = "<unknown>"
 
-        kwargs = {
-            "op": "function",
-            "name": sentry_span_name,
-            "origin": self.origin,
-        }
-
-        scope = sentry_sdk.get_current_scope()
-        parent_sentry_span = scope.span
-        if parent_sentry_span:
-            sentry_span = parent_sentry_span.start_child(**kwargs)
-        else:
-            sentry_span = scope.start_span(**kwargs)
+        span = sentry_sdk.start_span(op="function", name=sentry_span_name, origin=self.origin, only_if_parent=True)
+        span.__enter__()
 
         fields = metadata.get("fields", [])
         for field in fields:
             if self._include_tracing_fields():
-                sentry_span.set_data(field, attrs.get(field))
+                span.set_data(field, attrs.get(field))
             else:
-                sentry_span.set_data(field, SENSITIVE_DATA_SUBSTITUTE)
-
-        scope.span = sentry_span
-        return (parent_sentry_span, sentry_span)
+                span.set_data(field, SENSITIVE_DATA_SUBSTITUTE)
 
-    def on_close(self, span_id: str, span_state: TraceState) -> None:
-        if span_state is None:
-            return
+        return span
 
-        parent_sentry_span, sentry_span = span_state
-        sentry_span.finish()
-        sentry_sdk.get_current_scope().span = parent_sentry_span
+    def on_close(self, span_id: str, span: Optional[SentrySpan]) -> None:
+        if span is not None:
+            span.__exit__(None, None, None)
 
-    def on_record(self, span_id: str, values: str, span_state: TraceState) -> None:
-        if span_state is None:
-            return
-        _parent_sentry_span, sentry_span = span_state
 
-        deserialized_values = json.loads(values)
-        for key, value in deserialized_values.items():
-            if self._include_tracing_fields():
-                sentry_span.set_data(key, value)
-            else:
-                sentry_span.set_data(key, SENSITIVE_DATA_SUBSTITUTE)
+    def on_record(self, span_id: str, values: str, span: Optional[SentrySpan]) -> None:
+        if span is not None:
+            deserialized_values = json.loads(values)
+            for key, value in deserialized_values.items():
+                if self._include_tracing_fields():
+                    span.set_data(key, value)
+                else:
+                    span.set_data(key, SENSITIVE_DATA_SUBSTITUTE)
 
 
 class RustTracingIntegration(Integration):

sentry_sdk/tracing.py

@@ -1,3 +1,4 @@
+import json
 import uuid
 import random
 import time
@@ -1631,8 +1632,13 @@ def finish(self, end_timestamp=None):
 
     def to_json(self):
         # type: () -> dict[str, Any]
-        # TODO-neel-potel for sampling context
-        pass
+        """
+        Only meant for testing. Not used internally anymore.
+        """
+        if not isinstance(self._otel_span, ReadableSpan):
+            return {}
+        return json.loads(self._otel_span.to_json())
+
 
     def get_trace_context(self):
         # type: () -> dict[str, Any]

tests/integrations/rust_tracing/test_rust_tracing.py

@@ -11,7 +11,8 @@
     RustTracingLevel,
     EventTypeMapping,
 )
-from sentry_sdk import start_transaction, capture_message
+from sentry_sdk import start_span, capture_message
+from tests.conftest import ApproxDict
 
 
 def _test_event_type_mapping(metadata: Dict[str, object]) -> EventTypeMapping:
@@ -74,11 +75,11 @@ def test_on_new_span_on_close(sentry_init, capture_events):
     sentry_init(integrations=[integration], traces_sample_rate=1.0)
 
     events = capture_events()
-    with start_transaction():
+    with start_span():
         rust_tracing.new_span(RustTracingLevel.Info, 3)
 
         sentry_first_rust_span = sentry_sdk.get_current_span()
-        _, rust_first_rust_span = rust_tracing.spans[3]
+        rust_first_rust_span = rust_tracing.spans[3]
 
         assert sentry_first_rust_span == rust_first_rust_span
 
@@ -102,7 +103,7 @@ def test_on_new_span_on_close(sentry_init, capture_events):
     data = span["data"]
     assert data["use_memoized"]
     assert data["index"] == 10
-    assert data["version"] is None
+    assert "version" not in data
 
 
 def test_nested_on_new_span_on_close(sentry_init, capture_events):
@@ -115,23 +116,20 @@ def test_nested_on_new_span_on_close(sentry_init, capture_events):
     sentry_init(integrations=[integration], traces_sample_rate=1.0)
 
     events = capture_events()
-    with start_transaction():
+    with start_span():
         original_sentry_span = sentry_sdk.get_current_span()
 
         rust_tracing.new_span(RustTracingLevel.Info, 3, index_arg=10)
         sentry_first_rust_span = sentry_sdk.get_current_span()
-        _, rust_first_rust_span = rust_tracing.spans[3]
+        rust_first_rust_span = rust_tracing.spans[3]
 
         # Use a different `index_arg` value for the inner span to help
         # distinguish the two at the end of the test
         rust_tracing.new_span(RustTracingLevel.Info, 5, index_arg=9)
         sentry_second_rust_span = sentry_sdk.get_current_span()
-        rust_parent_span, rust_second_rust_span = rust_tracing.spans[5]
+        rust_second_rust_span = rust_tracing.spans[5]
 
         assert rust_second_rust_span == sentry_second_rust_span
-        assert rust_parent_span == sentry_first_rust_span
-        assert rust_parent_span == rust_first_rust_span
-        assert rust_parent_span != rust_second_rust_span
 
         rust_tracing.close_span(5)
 
@@ -171,12 +169,12 @@ def test_nested_on_new_span_on_close(sentry_init, capture_events):
     first_span_data = first_span["data"]
     assert first_span_data["use_memoized"]
     assert first_span_data["index"] == 10
-    assert first_span_data["version"] is None
+    assert "version" not in first_span_data
 
     second_span_data = second_span["data"]
     assert second_span_data["use_memoized"]
     assert second_span_data["index"] == 9
-    assert second_span_data["version"] is None
+    assert "version" not in second_span_data
 
 
 def test_on_new_span_without_transaction(sentry_init):
@@ -207,7 +205,7 @@ def test_on_event_exception(sentry_init, capture_events):
     events = capture_events()
     sentry_sdk.get_isolation_scope().clear_breadcrumbs()
 
-    with start_transaction():
+    with start_span():
         rust_tracing.new_span(RustTracingLevel.Info, 3)
 
         # Mapped to Exception
@@ -243,7 +241,7 @@ def test_on_event_breadcrumb(sentry_init, capture_events):
     events = capture_events()
     sentry_sdk.get_isolation_scope().clear_breadcrumbs()
 
-    with start_transaction():
+    with start_span():
         rust_tracing.new_span(RustTracingLevel.Info, 3)
 
         # Mapped to Breadcrumb
@@ -274,7 +272,7 @@ def test_on_event_event(sentry_init, capture_events):
     events = capture_events()
     sentry_sdk.get_isolation_scope().clear_breadcrumbs()
 
-    with start_transaction():
+    with start_span():
         rust_tracing.new_span(RustTracingLevel.Info, 3)
 
         # Mapped to Event
@@ -311,7 +309,7 @@ def test_on_event_ignored(sentry_init, capture_events):
     events = capture_events()
     sentry_sdk.get_isolation_scope().clear_breadcrumbs()
 
-    with start_transaction():
+    with start_span():
         rust_tracing.new_span(RustTracingLevel.Info, 3)
 
         # Ignored
@@ -344,7 +342,7 @@ def span_filter(metadata: Dict[str, object]) -> bool:
     sentry_init(integrations=[integration], traces_sample_rate=1.0)
 
     events = capture_events()
-    with start_transaction():
+    with start_span():
         original_sentry_span = sentry_sdk.get_current_span()
 
         # Span is not ignored
@@ -377,16 +375,16 @@ def test_record(sentry_init):
     )
     sentry_init(integrations=[integration], traces_sample_rate=1.0)
 
-    with start_transaction():
+    with start_span():
         rust_tracing.new_span(RustTracingLevel.Info, 3)
 
         span_before_record = sentry_sdk.get_current_span().to_json()
-        assert span_before_record["data"]["version"] is None
+        assert "version" not in span_before_record["attributes"]
 
         rust_tracing.record(3)
 
         span_after_record = sentry_sdk.get_current_span().to_json()
-        assert span_after_record["data"]["version"] == "memoized"
+        assert span_after_record["attributes"]["version"] == "memoized"
 
 
 def test_record_in_ignored_span(sentry_init):
@@ -403,18 +401,18 @@ def span_filter(metadata: Dict[str, object]) -> bool:
     )
     sentry_init(integrations=[integration], traces_sample_rate=1.0)
 
-    with start_transaction():
+    with start_span():
         rust_tracing.new_span(RustTracingLevel.Info, 3)
 
         span_before_record = sentry_sdk.get_current_span().to_json()
-        assert span_before_record["data"]["version"] is None
+        assert "version" not in span_before_record["attributes"]
 
         rust_tracing.new_span(RustTracingLevel.Trace, 5)
         rust_tracing.record(5)
 
         # `on_record()` should not do anything to the current Sentry span if the associated Rust span was ignored
         span_after_record = sentry_sdk.get_current_span().to_json()
-        assert span_after_record["data"]["version"] is None
+        assert "version" not in span_after_record["attributes"]
 
 
 @pytest.mark.parametrize(
@@ -443,33 +441,33 @@ def test_include_tracing_fields(
         traces_sample_rate=1.0,
         send_default_pii=send_default_pii,
     )
-    with start_transaction():
+    with start_span():
         rust_tracing.new_span(RustTracingLevel.Info, 3)
 
         span_before_record = sentry_sdk.get_current_span().to_json()
         if tracing_fields_expected:
-            assert span_before_record["data"]["version"] is None
+            assert "version" not in span_before_record["attributes"]
         else:
-            assert span_before_record["data"]["version"] == "[Filtered]"
+            assert span_before_record["attributes"]["version"] == "[Filtered]"
 
         rust_tracing.record(3)
 
         span_after_record = sentry_sdk.get_current_span().to_json()
 
         if tracing_fields_expected:
-            assert span_after_record["data"] == {
+            assert span_after_record["attributes"] == ApproxDict({
                 "thread.id": mock.ANY,
                 "thread.name": mock.ANY,
                 "use_memoized": True,
                 "version": "memoized",
                 "index": 10,
-            }
+            })
 
         else:
-            assert span_after_record["data"] == {
+            assert span_after_record["attributes"] == ApproxDict({
                 "thread.id": mock.ANY,
                 "thread.name": mock.ANY,
                 "use_memoized": "[Filtered]",
                 "version": "[Filtered]",
                 "index": "[Filtered]",
-            }
+            })