Profiler sampling context (#3840)

Author: sentrivana

sentry_sdk/integrations/opentelemetry/sampler.py

@@ -17,7 +17,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Optional, Sequence, Union
+    from typing import Any, Optional, Sequence, Union
     from opentelemetry.context import Context
     from opentelemetry.trace import Link, SpanKind
     from opentelemetry.trace.span import SpanContext
@@ -152,15 +152,9 @@ def should_sample(
         has_traces_sampler = callable(client.options.get("traces_sampler"))
 
         if is_root_span and has_traces_sampler:
-            sampling_context = {
-                "transaction_context": {
-                    "name": name,
-                    "op": attributes.get(SentrySpanAttribute.OP),
-                    "source": attributes.get(SentrySpanAttribute.SOURCE),
-                },
-                "parent_sampled": get_parent_sampled(parent_span_context, trace_id),
-            }
-            sampling_context.update(attributes)
+            sampling_context = create_sampling_context(
+                name, attributes, parent_span_context, trace_id
+            )
             sample_rate = client.options["traces_sampler"](sampling_context)
         else:
             # Check if there is a parent with a sampling decision
@@ -193,3 +187,19 @@ def should_sample(
 
     def get_description(self) -> str:
         return self.__class__.__name__
+
+
+def create_sampling_context(name, attributes, parent_span_context, trace_id):
+    # type: (str, Attributes, SpanContext, str) -> dict[str, Any]
+    sampling_context = {
+        "transaction_context": {
+            "name": name,
+            "op": attributes.get(SentrySpanAttribute.OP),
+            "source": attributes.get(SentrySpanAttribute.SOURCE),
+        },
+        "parent_sampled": get_parent_sampled(parent_span_context, trace_id),
+    }
+
+    sampling_context.update(attributes)
+
+    return sampling_context

sentry_sdk/integrations/opentelemetry/span_processor.py

@@ -20,6 +20,7 @@
     get_profiler_id,
 )
 from sentry_sdk.profiler.transaction_profiler import Profile
+from sentry_sdk.integrations.opentelemetry.sampler import create_sampling_context
 from sentry_sdk.integrations.opentelemetry.utils import (
     is_sentry_span,
     convert_from_otel_timestamp,
@@ -126,8 +127,10 @@ def _start_profile(self, span):
             # unix timestamp that is on span.start_time
             # setting it to 0 means the profiler will internally measure time on start
             profile = Profile(sampled, 0)
-            # TODO-neel-potel sampling context??
-            profile._set_initial_sampling_decision(sampling_context={})
+            sampling_context = create_sampling_context(
+                span.name, span.attributes, span.parent, span.context.trace_id
+            )
+            profile._set_initial_sampling_decision(sampling_context)
             profile.__enter__()
             set_sentry_meta(span, "profile", profile)
 

sentry_sdk/tracing_utils.py

@@ -33,7 +33,6 @@
     from typing import Generator
     from typing import Optional
     from typing import Union
-
     from types import FrameType
 
 
@@ -731,6 +730,3 @@ def get_current_span(scope=None):
     LOW_QUALITY_TRANSACTION_SOURCES,
     SENTRY_TRACE_HEADER_NAME,
 )
-
-if TYPE_CHECKING:
-    from sentry_sdk.tracing import Span

tests/tracing/test_sampling.py

@@ -307,3 +307,24 @@ def test_records_lost_event_only_if_traces_sampler_enabled(
 
     # Use Counter because order of calls does not matter
     assert Counter(record_lost_event_calls) == Counter(expected_record_lost_event_calls)
+
+
+@pytest.mark.parametrize("parent_sampling_decision", [True, False])
+def test_profiles_sampler_gets_sampling_context(sentry_init, parent_sampling_decision):
+    def dummy_profiles_sampler(sampling_context):
+        assert sampling_context["transaction_context"] == {
+            "name": "dogpark",
+            "op": "op",
+            "source": "custom",
+        }
+        assert sampling_context["parent_sampled"] == parent_sampling_decision
+        return 1.0
+
+    sentry_init(traces_sample_rate=1.0, profiles_sampler=dummy_profiles_sampler)
+
+    sentry_trace = "12312012123120121231201212312012-1121201211212012-{}".format(
+        int(parent_sampling_decision)
+    )
+    with sentry_sdk.continue_trace({"sentry-trace": sentry_trace}):
+        with sentry_sdk.start_span(name="dogpark", op="op"):
+            pass