From 8d2a2013525d8c4db30d7e6d1593ea207f80770e Mon Sep 17 00:00:00 2001 From: Neel Shah Date: Wed, 25 Sep 2024 13:15:05 -0400 Subject: [PATCH] Head SDK DSC population Note that transaction name is missing here for now --- .../integrations/opentelemetry/consts.py | 4 +- .../integrations/opentelemetry/sampler.py | 53 ++++++++++----- .../integrations/opentelemetry/scope.py | 1 - .../integrations/opentelemetry/utils.py | 65 +++++++++++++++++-- sentry_sdk/tracing.py | 36 +++++++--- 5 files changed, 125 insertions(+), 34 deletions(-) diff --git a/sentry_sdk/integrations/opentelemetry/consts.py b/sentry_sdk/integrations/opentelemetry/consts.py index cb088f13a5..6409d2822d 100644 --- a/sentry_sdk/integrations/opentelemetry/consts.py +++ b/sentry_sdk/integrations/opentelemetry/consts.py @@ -1,4 +1,5 @@ from opentelemetry.context import create_key +from sentry_sdk.tracing_utils import Baggage # propagation keys @@ -11,7 +12,8 @@ SENTRY_USE_CURRENT_SCOPE_KEY = create_key("sentry_use_current_scope") SENTRY_USE_ISOLATION_SCOPE_KEY = create_key("sentry_use_isolation_scope") -SENTRY_TRACE_STATE_DROPPED = "sentry_dropped" +TRACESTATE_SAMPLED_KEY = Baggage.SENTRY_PREFIX + "sampled" +TRACESTATE_SAMPLE_RATE_KEY = Baggage.SENTRY_PREFIX + "sample_rate" OTEL_SENTRY_CONTEXT = "otel" SPAN_ORIGIN = "auto.otel" diff --git a/sentry_sdk/integrations/opentelemetry/sampler.py b/sentry_sdk/integrations/opentelemetry/sampler.py index 5fa7c9e1e8..404957f028 100644 --- a/sentry_sdk/integrations/opentelemetry/sampler.py +++ b/sentry_sdk/integrations/opentelemetry/sampler.py @@ -1,3 +1,4 @@ +from typing import cast from random import random from opentelemetry import trace @@ -6,13 +7,17 @@ from opentelemetry.trace.span import TraceState import sentry_sdk -from sentry_sdk.integrations.opentelemetry.consts import SENTRY_TRACE_STATE_DROPPED from sentry_sdk.tracing_utils import has_tracing_enabled from sentry_sdk.utils import is_valid_sample_rate, logger +from sentry_sdk.integrations.opentelemetry.consts import ( + TRACESTATE_SAMPLED_KEY, + TRACESTATE_SAMPLE_RATE_KEY, +) -from typing import TYPE_CHECKING, Optional, Sequence +from typing import TYPE_CHECKING if TYPE_CHECKING: + from typing import Optional, Sequence, Union from opentelemetry.context import Context from opentelemetry.trace import Link, SpanKind from opentelemetry.trace.span import SpanContext @@ -32,30 +37,42 @@ def get_parent_sampled(parent_context, trace_id): if parent_context.trace_flags.sampled: return True - dropped = parent_context.trace_state.get(SENTRY_TRACE_STATE_DROPPED) == "true" - if dropped: + dsc_sampled = parent_context.trace_state.get(TRACESTATE_SAMPLED_KEY) + if dsc_sampled == "true": + return True + elif dsc_sampled == "false": return False - # TODO-anton: fall back to sampling decision in DSC (for this die DSC needs to be set in the trace_state) - return None -def dropped_result(span_context): - # type: (SpanContext) -> SamplingResult - trace_state = span_context.trace_state.update(SENTRY_TRACE_STATE_DROPPED, "true") +def dropped_result(span_context, attributes, sample_rate=None): + # type: (SpanContext, Attributes, Optional[float]) -> SamplingResult + # note that trace_state.add will NOT overwrite existing entries + # so these will only be added the first time in a root span sampling decision + trace_state = span_context.trace_state.add(TRACESTATE_SAMPLED_KEY, "false") + if sample_rate: + trace_state = trace_state.add(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate)) return SamplingResult( Decision.DROP, + attributes=attributes, trace_state=trace_state, ) -def sampled_result(span_context): - # type: (SpanContext) -> SamplingResult +def sampled_result(span_context, attributes, sample_rate): + # type: (SpanContext, Attributes, float) -> SamplingResult + # note that trace_state.add will NOT overwrite existing entries + # so these will only be added the first time in a root span sampling decision + trace_state = span_context.trace_state.add(TRACESTATE_SAMPLED_KEY, "true").add( + TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate) + ) + return SamplingResult( Decision.RECORD_AND_SAMPLE, - trace_state=span_context.trace_state, + attributes=attributes, + trace_state=trace_state, ) @@ -77,7 +94,7 @@ def should_sample( # No tracing enabled, thus no sampling if not has_tracing_enabled(client.options): - return dropped_result(parent_span_context) + return dropped_result(parent_span_context, attributes) sample_rate = None @@ -112,16 +129,16 @@ def should_sample( logger.warning( f"[Tracing] Discarding {name} because of invalid sample rate." ) - return dropped_result(parent_span_context) + return dropped_result(parent_span_context, attributes) # Roll the dice on sample rate - sampled = random() < float(sample_rate) + sample_rate = float(cast("Union[bool, float, int]", sample_rate)) + sampled = random() < sample_rate - # TODO-neel-potel set sample rate as attribute for DSC if sampled: - return sampled_result(parent_span_context) + return sampled_result(parent_span_context, attributes, sample_rate) else: - return dropped_result(parent_span_context) + return dropped_result(parent_span_context, attributes, sample_rate) def get_description(self) -> str: return self.__class__.__name__ diff --git a/sentry_sdk/integrations/opentelemetry/scope.py b/sentry_sdk/integrations/opentelemetry/scope.py index fc76c4ffdb..11714fda53 100644 --- a/sentry_sdk/integrations/opentelemetry/scope.py +++ b/sentry_sdk/integrations/opentelemetry/scope.py @@ -8,7 +8,6 @@ SpanContext, NonRecordingSpan, TraceFlags, - TraceState, use_span, ) diff --git a/sentry_sdk/integrations/opentelemetry/utils.py b/sentry_sdk/integrations/opentelemetry/utils.py index 982ecb4509..fcfec97f5c 100644 --- a/sentry_sdk/integrations/opentelemetry/utils.py +++ b/sentry_sdk/integrations/opentelemetry/utils.py @@ -15,6 +15,7 @@ from opentelemetry.semconv.trace import SpanAttributes from opentelemetry.sdk.trace import ReadableSpan +import sentry_sdk from sentry_sdk.utils import Dsn from sentry_sdk.consts import SPANSTATUS from sentry_sdk.tracing import get_span_status_from_http_code, DEFAULT_SPAN_ORIGIN @@ -24,7 +25,7 @@ from sentry_sdk._types import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Optional, Mapping, Sequence, Union, ItemsView + from typing import Any, Optional, Mapping, Sequence, Union from sentry_sdk._types import OtelExtractedSpanData @@ -300,9 +301,8 @@ def get_trace_context(span, span_data=None): if span.attributes: trace_context["data"] = dict(span.attributes) - trace_context["dynamic_sampling_context"] = dsc_from_trace_state( - span.context.trace_state - ) + trace_state = get_trace_state(span) + trace_context["dynamic_sampling_context"] = dsc_from_trace_state(trace_state) # TODO-neel-potel profiler thread_id, thread_name @@ -319,6 +319,11 @@ def trace_state_from_baggage(baggage): return TraceState(items) +def baggage_from_trace_state(trace_state): + # type: (TraceState) -> Baggage + return Baggage(dsc_from_trace_state(trace_state)) + + def serialize_trace_state(trace_state): # type: (TraceState) -> str sentry_items = [] @@ -336,3 +341,55 @@ def dsc_from_trace_state(trace_state): key = re.sub(Baggage.SENTRY_PREFIX_REGEX, "", k) dsc[key] = v return dsc + + +def has_incoming_trace(trace_state): + # type: (TraceState) -> bool + """ + The existence a sentry-trace_id in the baggage implies we continued an upstream trace. + """ + return (Baggage.SENTRY_PREFIX + "trace_id") in trace_state + + +def get_trace_state(span): + # type: (Union[Span, ReadableSpan]) -> TraceState + """ + Get the existing trace_state with sentry items + or populate it if we are the head SDK. + """ + span_context = span.get_span_context() + if not span_context: + return TraceState() + + trace_state = span_context.trace_state + + if has_incoming_trace(trace_state): + return trace_state + else: + client = sentry_sdk.get_client() + if not client.is_active(): + return trace_state + + options = client.options or {} + + trace_state = trace_state.update( + Baggage.SENTRY_PREFIX + "trace_id", format_trace_id(span_context.trace_id) + ) + + if options.get("environment"): + trace_state = trace_state.update( + Baggage.SENTRY_PREFIX + "environment", options["environment"] + ) + + if options.get("release"): + trace_state = trace_state.update( + Baggage.SENTRY_PREFIX + "release", options["release"] + ) + + if options.get("dsn"): + trace_state = trace_state.update( + Baggage.SENTRY_PREFIX + "public_key", Dsn(options["dsn"]).public_key + ) + + # TODO-neel-potel head dsc transaction name + return trace_state diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py index 46c19996bd..1c3cb5b3f0 100644 --- a/sentry_sdk/tracing.py +++ b/sentry_sdk/tracing.py @@ -5,7 +5,12 @@ from datetime import datetime, timedelta, timezone from opentelemetry import trace as otel_trace, context -from opentelemetry.trace import format_trace_id, format_span_id, Span as OtelSpan +from opentelemetry.trace import ( + format_trace_id, + format_span_id, + Span as OtelSpan, + TraceState, +) from opentelemetry.trace.status import StatusCode from opentelemetry.sdk.trace import ReadableSpan @@ -1453,8 +1458,7 @@ def iter_headers(self): serialize_trace_state, ) - trace_state = self._otel_span.get_span_context().trace_state - yield BAGGAGE_HEADER_NAME, serialize_trace_state(trace_state) + yield BAGGAGE_HEADER_NAME, serialize_trace_state(self.trace_state) def to_traceparent(self): # type: () -> str @@ -1471,10 +1475,26 @@ def to_traceparent(self): return traceparent + @property + def trace_state(self): + # type: () -> TraceState + from sentry_sdk.integrations.opentelemetry.utils import ( + get_trace_state, + ) + + return get_trace_state(self._otel_span) + def to_baggage(self): - # type: () -> Optional[Baggage] - # TODO-neel-potel head SDK populate baggage mess - pass + # type: () -> Baggage + return self.get_baggage() + + def get_baggage(self): + # type: () -> Baggage + from sentry_sdk.integrations.opentelemetry.utils import ( + baggage_from_trace_state, + ) + + return baggage_from_trace_state(self.trace_state) def set_tag(self, key, value): # type: (str, Any) -> None @@ -1568,10 +1588,6 @@ def set_context(self, key, value): # type: (str, Any) -> None pass - def get_baggage(self): - # type: () -> Baggage - pass - if TYPE_CHECKING: