Fix deferred tracing decisions for twp

Author: sl0thentr0py

sentry_sdk/integrations/opentelemetry/sampler.py

@@ -34,10 +34,14 @@ def get_parent_sampled(parent_context, trace_id):
     # Only inherit sample rate if `traceId` is the same
     if is_span_context_valid and parent_context.trace_id == trace_id:
         # this is getSamplingDecision in JS
+        # if there was no sampling flag, defer the decision
+        dsc_sampled = parent_context.trace_state.get(TRACESTATE_SAMPLED_KEY)
+        if dsc_sampled == "deferred":
+            return None
+
         if parent_context.trace_flags.sampled is not None:
             return parent_context.trace_flags.sampled
 
-        dsc_sampled = parent_context.trace_state.get(TRACESTATE_SAMPLED_KEY)
         if dsc_sampled == "true":
             return True
         elif dsc_sampled == "false":
@@ -53,6 +57,8 @@ def dropped_result(parent_span_context, attributes, sample_rate=None):
 
     if TRACESTATE_SAMPLED_KEY not in trace_state:
         trace_state = trace_state.add(TRACESTATE_SAMPLED_KEY, "false")
+    elif trace_state.get(TRACESTATE_SAMPLED_KEY) == "deferred":
+        trace_state = trace_state.update(TRACESTATE_SAMPLED_KEY, "false")
 
     if sample_rate and TRACESTATE_SAMPLE_RATE_KEY not in trace_state:
         trace_state = trace_state.add(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate))
@@ -88,6 +94,9 @@ def sampled_result(span_context, attributes, sample_rate):
 
     if TRACESTATE_SAMPLED_KEY not in trace_state:
         trace_state = trace_state.add(TRACESTATE_SAMPLED_KEY, "true")
+    elif trace_state.get(TRACESTATE_SAMPLED_KEY) == "deferred":
+        trace_state = trace_state.update(TRACESTATE_SAMPLED_KEY, "true")
+
     if TRACESTATE_SAMPLE_RATE_KEY not in trace_state:
         trace_state = trace_state.add(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate))
 

sentry_sdk/integrations/opentelemetry/scope.py

@@ -6,6 +6,7 @@
     SpanContext,
     NonRecordingSpan,
     TraceFlags,
+    TraceState,
     use_span,
 )
 
@@ -14,6 +15,7 @@
     SENTRY_FORK_ISOLATION_SCOPE_KEY,
     SENTRY_USE_CURRENT_SCOPE_KEY,
     SENTRY_USE_ISOLATION_SCOPE_KEY,
+    TRACESTATE_SAMPLED_KEY,
 )
 from sentry_sdk.integrations.opentelemetry.utils import trace_state_from_baggage
 from sentry_sdk.scope import Scope, ScopeType
@@ -96,10 +98,15 @@ def _incoming_otel_span_context(self):
             else TraceFlags.DEFAULT
         )
 
-        # TODO-neel-potel do we need parent and sampled like JS?
-        trace_state = None
         if self._propagation_context.baggage:
             trace_state = trace_state_from_baggage(self._propagation_context.baggage)
+        else:
+            trace_state = TraceState()
+
+        # for twp to work, we also need to consider deferred sampling when the sampling
+        # flag is not present, so the above TraceFlags are not sufficient
+        if self._propagation_context.parent_sampled is None:
+            trace_state = trace_state.add(TRACESTATE_SAMPLED_KEY, "deferred")
 
         span_context = SpanContext(
             trace_id=int(self._propagation_context.trace_id, 16),  # type: ignore

sentry_sdk/tracing.py

@@ -1311,6 +1311,8 @@ def __exit__(self, ty, value, tb):
         # type: (Optional[Any], Optional[Any], Optional[Any]) -> None
         if value is not None:
             self.set_status(SPANSTATUS.INTERNAL_ERROR)
+        else:
+            self.set_status(SPANSTATUS.OK)
 
         self.finish()
         context.detach(self._ctx_token)

tests/integrations/opentelemetry/test_compat.py

@@ -9,7 +9,7 @@ def test_transaction_name_span_description_compat(
 
     events = capture_events()
 
-    with sentry_sdk.start_transaction(
+    with sentry_sdk.start_span(
         name="trx-name",
         op="trx-op",
     ) as trx:
@@ -33,13 +33,12 @@ def test_transaction_name_span_description_compat(
     assert spn.__class__.__name__ == "POTelSpan"
     assert spn.op == "span-op"
     assert spn.description == "span-desc"
-    assert spn.name is None
+    assert spn.name == "span-desc"
 
     assert spn._otel_span is not None
     assert spn._otel_span.name == "span-desc"
     assert spn._otel_span.attributes["sentry.op"] == "span-op"
     assert spn._otel_span.attributes["sentry.description"] == "span-desc"
-    assert "sentry.name" not in spn._otel_span.attributes
 
     transaction = events[0]
     assert transaction["transaction"] == "trx-name"
@@ -53,4 +52,3 @@ def test_transaction_name_span_description_compat(
     assert span["op"] == "span-op"
     assert span["data"]["sentry.op"] == "span-op"
     assert span["data"]["sentry.description"] == "span-desc"
-    assert "sentry.name" not in span["data"]

tests/integrations/opentelemetry/test_experimental.py

@@ -1,8 +1,8 @@
 import pytest
-
 from opentelemetry import trace
 
 import sentry_sdk
+from tests.conftest import ApproxDict
 
 
 tracer = trace.get_tracer(__name__)
@@ -43,7 +43,6 @@ def test_root_span_transaction_payload_started_with_otel_only(capture_envelopes)
     assert "span_id" in trace_context
     assert trace_context["origin"] == "manual"
     assert trace_context["op"] == "request"
-    assert trace_context["status"] == "ok"
 
     assert payload["spans"] == []
 
@@ -63,7 +62,6 @@ def test_child_span_payload_started_with_otel_only(capture_envelopes):
     assert span["op"] == "db"
     assert span["description"] == "db"
     assert span["origin"] == "manual"
-    assert span["status"] == "ok"
     assert span["span_id"] is not None
     assert span["trace_id"] == payload["contexts"]["trace"]["trace_id"]
     assert span["parent_span_id"] == payload["contexts"]["trace"]["span_id"]
@@ -222,8 +220,8 @@ def test_span_attributes_in_data_started_with_otel(capture_envelopes):
     (item,) = envelope.items
     payload = item.payload.json
 
-    assert payload["contexts"]["trace"]["data"] == {"foo": "bar", "baz": 42}
-    assert payload["spans"][0]["data"] == {"abc": 99, "def": "moo"}
+    assert payload["contexts"]["trace"]["data"] == ApproxDict({"foo": "bar", "baz": 42})
+    assert payload["spans"][0]["data"] == ApproxDict({"abc": 99, "def": "moo"})
 
 
 def test_span_data_started_with_sentry(capture_envelopes):
@@ -238,18 +236,22 @@ def test_span_data_started_with_sentry(capture_envelopes):
     (item,) = envelope.items
     payload = item.payload.json
 
-    assert payload["contexts"]["trace"]["data"] == {
-        "foo": "bar",
-        "sentry.origin": "manual",
-        "sentry.description": "request",
-        "sentry.op": "http",
-    }
-    assert payload["spans"][0]["data"] == {
-        "baz": 42,
-        "sentry.origin": "manual",
-        "sentry.description": "statement",
-        "sentry.op": "db",
-    }
+    assert payload["contexts"]["trace"]["data"] == ApproxDict(
+        {
+            "foo": "bar",
+            "sentry.origin": "manual",
+            "sentry.description": "request",
+            "sentry.op": "http",
+        }
+    )
+    assert payload["spans"][0]["data"] == ApproxDict(
+        {
+            "baz": 42,
+            "sentry.origin": "manual",
+            "sentry.description": "statement",
+            "sentry.op": "db",
+        }
+    )
 
 
 def test_transaction_tags_started_with_otel(capture_envelopes):

tests/integrations/opentelemetry/test_potel.py

@@ -71,17 +71,13 @@ def test_sampling_traces_sample_rate_50(sentry_init, capture_envelopes):
 
     envelopes = capture_envelopes()
 
-    with mock.patch(
-        "sentry_sdk.integrations.opentelemetry.sampler.random", return_value=0.2
-    ):  # drop
+    with mock.patch("random.random", return_value=0.2):  # drop
         with sentry_sdk.start_span(description="request a"):
             with sentry_sdk.start_span(description="cache a"):
                 with sentry_sdk.start_span(description="db a"):
                     ...
 
-    with mock.patch(
-        "sentry_sdk.integrations.opentelemetry.sampler.random", return_value=0.7
-    ):  # keep
+    with mock.patch("random.random", return_value=0.7):  # keep
         with sentry_sdk.start_span(description="request b"):
             with sentry_sdk.start_span(description="cache b"):
                 with sentry_sdk.start_span(description="db b"):
@@ -101,46 +97,34 @@ def test_sampling_traces_sample_rate_50(sentry_init, capture_envelopes):
 def test_sampling_traces_sampler(sentry_init, capture_envelopes):
     def keep_only_a(sampling_context):
         if " a" in sampling_context["transaction_context"]["name"]:
-            return 0.05
+            return 1
         else:
             return 0
 
-    sentry_init(
-        traces_sample_rate=1.0,
-        traces_sampler=keep_only_a,
-    )
+    sentry_init(traces_sampler=keep_only_a)
 
     envelopes = capture_envelopes()
 
-    # Make sure random() always returns the same values
-    with mock.patch(
-        "sentry_sdk.integrations.opentelemetry.sampler.random",
-        side_effect=[0.04 for _ in range(12)],
-    ):
-
-        with sentry_sdk.start_span(description="request a"):  # keep
-            with sentry_sdk.start_span(description="cache a"):  # keep
-                with sentry_sdk.start_span(description="db a"):  # keep
-                    ...
+    # children inherit from root spans
+    with sentry_sdk.start_span(description="request a"):  # keep
+        with sentry_sdk.start_span(description="cache a"):
+            with sentry_sdk.start_span(description="db a"):
+                ...
 
-        with sentry_sdk.start_span(description="request b"):  # drop
-            with sentry_sdk.start_span(description="cache b"):  # drop
-                with sentry_sdk.start_span(description="db b"):  # drop
-                    ...
+    with sentry_sdk.start_span(description="request b"):  # drop
+        with sentry_sdk.start_span(description="cache b"):
+            with sentry_sdk.start_span(description="db b"):
+                ...
 
-        with sentry_sdk.start_span(description="request c"):  # drop
-            with sentry_sdk.start_span(
-                description="cache a c"
-            ):  # keep (but trx dropped, so not collected)
-                with sentry_sdk.start_span(
-                    description="db a c"
-                ):  # keep (but trx dropped, so not collected)
-                    ...
+    with sentry_sdk.start_span(description="request c"):  # drop
+        with sentry_sdk.start_span(description="cache a c"):
+            with sentry_sdk.start_span(description="db a c"):
+                ...
 
-        with sentry_sdk.start_span(description="new a c"):  # keep
-            with sentry_sdk.start_span(description="cache c"):  # drop
-                with sentry_sdk.start_span(description="db c"):  # drop
-                    ...
+    with sentry_sdk.start_span(description="new a c"):  # keep
+        with sentry_sdk.start_span(description="cache c"):
+            with sentry_sdk.start_span(description="db c"):
+                ...
 
     assert len(envelopes) == 2
     (envelope1, envelope2) = envelopes
@@ -150,7 +134,7 @@ def keep_only_a(sampling_context):
     assert transaction1["transaction"] == "request a"
     assert len(transaction1["spans"]) == 2
     assert transaction2["transaction"] == "new a c"
-    assert len(transaction2["spans"]) == 0
+    assert len(transaction2["spans"]) == 2
 
 
 def test_sampling_traces_sampler_boolean(sentry_init, capture_envelopes):
@@ -168,21 +152,21 @@ def keep_only_a(sampling_context):
     envelopes = capture_envelopes()
 
     with sentry_sdk.start_span(description="request a"):  # keep
-        with sentry_sdk.start_span(description="cache a"):  # keep
-            with sentry_sdk.start_span(description="db X"):  # drop
+        with sentry_sdk.start_span(description="cache a"):
+            with sentry_sdk.start_span(description="db X"):
                 ...
 
     with sentry_sdk.start_span(description="request b"):  # drop
-        with sentry_sdk.start_span(description="cache b"):  # drop
-            with sentry_sdk.start_span(description="db b"):  # drop
+        with sentry_sdk.start_span(description="cache b"):
+            with sentry_sdk.start_span(description="db b"):
                 ...
 
     assert len(envelopes) == 1
     (envelope,) = envelopes
     transaction = envelope.items[0].payload.json
 
     assert transaction["transaction"] == "request a"
-    assert len(transaction["spans"]) == 1
+    assert len(transaction["spans"]) == 2
 
 
 @pytest.mark.parametrize(
@@ -237,21 +221,24 @@ def test_sampling_parent_sampled(
 
 
 @pytest.mark.parametrize(
-    "traces_sample_rate, expected_num_of_envelopes",
+    "traces_sample_rate, upstream_sampled, expected_num_of_envelopes",
     [
         # special case for testing, do not pass any traces_sample_rate to init() (the default traces_sample_rate=None will be used)
-        (-1, 0),
+        (-1, 0, 0),
         # traces_sample_rate=None means do not create new traces, and also do not continue incoming traces. So, no envelopes at all.
-        (None, 0),
+        (None, 1, 0),
         # traces_sample_rate=0 means do not create new traces (0% of the requests), but continue incoming traces. So envelopes will be created only if there is an incoming trace.
-        (0, 0),
+        (0, 0, 0),
+        (0, 1, 1),
         # traces_sample_rate=1 means create new traces for 100% of requests (and also continue incoming traces, of course).
-        (1, 1),
+        (1, 0, 0),
+        (1, 1, 1),
     ],
 )
 def test_sampling_parent_dropped(
     sentry_init,
     traces_sample_rate,
+    upstream_sampled,
     expected_num_of_envelopes,
     capture_envelopes,
 ):
@@ -265,7 +252,7 @@ def test_sampling_parent_dropped(
 
     # The upstream service has dropped the request
     headers = {
-        "sentry-trace": "771a43a4192642f0b136d5159a501700-1234567890abcdef-0",
+        "sentry-trace": f"771a43a4192642f0b136d5159a501700-1234567890abcdef-{upstream_sampled}",
     }
     with sentry_sdk.continue_trace(headers):
         with sentry_sdk.start_span(description="request a"):