From 39c2ed2385b12adfc371118edc44d48df06131ef Mon Sep 17 00:00:00 2001 From: Neel Shah Date: Thu, 12 Sep 2024 15:57:19 +0200 Subject: [PATCH] Make DSC work for outgoing traces * Add parsed baggage sentry items as items on the `sampling_context.trace_state` in `continue_trace` * Fix sampler to propagate the `trace_state` to children in both sampled and dropped cases * make `iter_headers` work * make `get_trace_context` work * add `dynamic_sampling_context` in `trace_context` so that it can be picked up by the client and added in the envelope header --- sentry_sdk/_types.py | 2 + .../opentelemetry/potel_span_processor.py | 23 ++--- .../integrations/opentelemetry/sampler.py | 37 ++++---- .../integrations/opentelemetry/scope.py | 13 +-- .../integrations/opentelemetry/utils.py | 84 +++++++++++++++++-- sentry_sdk/tracing.py | 23 ++++- 6 files changed, 130 insertions(+), 52 deletions(-) diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py index 0ee3921862..19446f0f86 100644 --- a/sentry_sdk/_types.py +++ b/sentry_sdk/_types.py @@ -197,3 +197,5 @@ class SDKInfo(TypedDict): ) HttpStatusCodeRange = Union[int, Container[int]] + + OtelExtractedSpanData = tuple[str, str, Optional[str], Optional[int], Optional[str]] diff --git a/sentry_sdk/integrations/opentelemetry/potel_span_processor.py b/sentry_sdk/integrations/opentelemetry/potel_span_processor.py index d61b5f8782..63a9acb9db 100644 --- a/sentry_sdk/integrations/opentelemetry/potel_span_processor.py +++ b/sentry_sdk/integrations/opentelemetry/potel_span_processor.py @@ -18,6 +18,7 @@ convert_from_otel_timestamp, extract_span_attributes, extract_span_data, + get_trace_context, ) from sentry_sdk.integrations.opentelemetry.consts import ( OTEL_SENTRY_CONTEXT, @@ -136,26 +137,12 @@ def _root_span_to_transaction_event(self, span): if event is None: return None - trace_id = format_trace_id(span.context.trace_id) - span_id = format_span_id(span.context.span_id) - parent_span_id = format_span_id(span.parent.span_id) if span.parent else None - - (op, description, status, _, origin) = extract_span_data(span) - - trace_context = { - "trace_id": trace_id, - "span_id": span_id, - "origin": origin or DEFAULT_SPAN_ORIGIN, - "op": op, - "status": status, - } # type: dict[str, Any] - - if parent_span_id: - trace_context["parent_span_id"] = parent_span_id - if span.attributes: - trace_context["data"] = dict(span.attributes) + span_data = extract_span_data(span) + (_, description, _, _, _) = span_data + trace_context = get_trace_context(span, span_data=span_data) contexts = {"trace": trace_context} + if span.resource.attributes: contexts[OTEL_SENTRY_CONTEXT] = {"resource": dict(span.resource.attributes)} diff --git a/sentry_sdk/integrations/opentelemetry/sampler.py b/sentry_sdk/integrations/opentelemetry/sampler.py index 445c2edd02..5fa7c9e1e8 100644 --- a/sentry_sdk/integrations/opentelemetry/sampler.py +++ b/sentry_sdk/integrations/opentelemetry/sampler.py @@ -41,16 +41,21 @@ def get_parent_sampled(parent_context, trace_id): return None -def dropped(parent_context=None): - # type: (Optional[SpanContext]) -> SamplingResult - trace_state = parent_context.trace_state if parent_context is not None else None - updated_trace_context = trace_state or TraceState() - updated_trace_context = updated_trace_context.update( - SENTRY_TRACE_STATE_DROPPED, "true" - ) +def dropped_result(span_context): + # type: (SpanContext) -> SamplingResult + trace_state = span_context.trace_state.update(SENTRY_TRACE_STATE_DROPPED, "true") + return SamplingResult( Decision.DROP, - trace_state=updated_trace_context, + trace_state=trace_state, + ) + + +def sampled_result(span_context): + # type: (SpanContext) -> SamplingResult + return SamplingResult( + Decision.RECORD_AND_SAMPLE, + trace_state=span_context.trace_state, ) @@ -68,12 +73,11 @@ def should_sample( # type: (...) -> SamplingResult client = sentry_sdk.get_client() - parent_span = trace.get_current_span(parent_context) - parent_context = parent_span.get_span_context() if parent_span else None + parent_span_context = trace.get_current_span(parent_context).get_span_context() # No tracing enabled, thus no sampling if not has_tracing_enabled(client.options): - return dropped(parent_context) + return dropped_result(parent_span_context) sample_rate = None @@ -89,14 +93,14 @@ def should_sample( "transaction_context": { "name": name, }, - "parent_sampled": get_parent_sampled(parent_context, trace_id), + "parent_sampled": get_parent_sampled(parent_span_context, trace_id), } sample_rate = client.options["traces_sampler"](sampling_context) else: # Check if there is a parent with a sampling decision - parent_sampled = get_parent_sampled(parent_context, trace_id) + parent_sampled = get_parent_sampled(parent_span_context, trace_id) if parent_sampled is not None: sample_rate = parent_sampled else: @@ -108,15 +112,16 @@ def should_sample( logger.warning( f"[Tracing] Discarding {name} because of invalid sample rate." ) - return dropped(parent_context) + return dropped_result(parent_span_context) # Roll the dice on sample rate sampled = random() < float(sample_rate) + # TODO-neel-potel set sample rate as attribute for DSC if sampled: - return SamplingResult(Decision.RECORD_AND_SAMPLE) + return sampled_result(parent_span_context) else: - return dropped(parent_context) + return dropped_result(parent_span_context) 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 a3eb1f3268..8d8c977b35 100644 --- a/sentry_sdk/integrations/opentelemetry/scope.py +++ b/sentry_sdk/integrations/opentelemetry/scope.py @@ -4,7 +4,7 @@ from contextlib import contextmanager from opentelemetry.context import get_value, set_value, attach, detach, get_current -from opentelemetry.trace import SpanContext, NonRecordingSpan, TraceFlags, use_span +from opentelemetry.trace import SpanContext, NonRecordingSpan, TraceFlags, TraceState, use_span from sentry_sdk.integrations.opentelemetry.consts import ( SENTRY_SCOPES_KEY, @@ -12,6 +12,7 @@ SENTRY_USE_CURRENT_SCOPE_KEY, SENTRY_USE_ISOLATION_SCOPE_KEY, ) +from sentry_sdk.integrations.opentelemetry.utils import trace_state_from_baggage from sentry_sdk.scope import Scope, ScopeType from sentry_sdk.tracing import POTelSpan from sentry_sdk._types import TYPE_CHECKING @@ -93,15 +94,17 @@ def _incoming_otel_span_context(self): else TraceFlags.DEFAULT ) - # TODO-neel-potel tracestate + # 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) + span_context = SpanContext( trace_id=int(self._propagation_context.trace_id, 16), # type: ignore span_id=int(self._propagation_context.parent_span_id, 16), # type: ignore is_remote=True, trace_flags=trace_flags, - # TODO-anton: add trace_state (mapping[str,str]) with the parentSpanId, dsc and sampled from self._propagation_context - # trace_state={ - # } + trace_state=trace_state, ) return span_context diff --git a/sentry_sdk/integrations/opentelemetry/utils.py b/sentry_sdk/integrations/opentelemetry/utils.py index 2d16a0835d..571389fb60 100644 --- a/sentry_sdk/integrations/opentelemetry/utils.py +++ b/sentry_sdk/integrations/opentelemetry/utils.py @@ -1,20 +1,24 @@ +import re from typing import cast from datetime import datetime, timezone -from opentelemetry.trace import SpanKind, StatusCode +from urllib3.util import parse_url as urlparse +from urllib.parse import quote +from opentelemetry.trace import Span, SpanKind, StatusCode, format_trace_id, format_span_id, TraceState from opentelemetry.semconv.trace import SpanAttributes from opentelemetry.sdk.trace import ReadableSpan -from sentry_sdk.consts import SPANSTATUS -from sentry_sdk.tracing import get_span_status_from_http_code -from sentry_sdk.integrations.opentelemetry.consts import SentrySpanAttribute -from urllib3.util import parse_url as urlparse 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 +from sentry_sdk.tracing_utils import Baggage +from sentry_sdk.integrations.opentelemetry.consts import SentrySpanAttribute from sentry_sdk._types import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Optional, Mapping, Sequence, Union + from typing import Any, Optional, Mapping, Sequence, Union, ItemsView + from sentry_sdk._types import OtelExtractedSpanData GRPC_ERROR_MAP = { @@ -87,7 +91,7 @@ def convert_to_otel_timestamp(time): def extract_span_data(span): - # type: (ReadableSpan) -> tuple[str, str, Optional[str], Optional[int], Optional[str]] + # type: (ReadableSpan) -> OtelExtractedSpanData op = span.name description = span.name status, http_status = extract_span_status(span) @@ -125,7 +129,7 @@ def extract_span_data(span): def span_data_for_http_method(span): - # type: (ReadableSpan) -> tuple[str, str, Optional[str], Optional[int], Optional[str]] + # type: (ReadableSpan) -> OtelExtractedSpanData span_attributes = span.attributes or {} op = "http" @@ -167,7 +171,7 @@ def span_data_for_http_method(span): def span_data_for_db_query(span): - # type: (ReadableSpan) -> tuple[str, str, Optional[str], Optional[int], Optional[str]] + # type: (ReadableSpan) -> OtelExtractedSpanData span_attributes = span.attributes or {} op = "db" @@ -261,3 +265,65 @@ def extract_span_attributes(span, namespace): extracted_attrs[key] = value return extracted_attrs + + +def get_trace_context(span, span_data=None): + # type: (ReadableSpan, Optional[OtelExtractedSpanData]) -> dict[str, Any] + if not span.context: + return {} + + trace_id = format_trace_id(span.context.trace_id) + span_id = format_span_id(span.context.span_id) + parent_span_id = format_span_id(span.parent.span_id) if span.parent else None + + if span_data is None: + span_data = extract_span_data(span) + + (op, _, status, _, origin) = span_data + + trace_context = { + "trace_id": trace_id, + "span_id": span_id, + "parent_span_id": parent_span_id, + "op": op, + "origin": origin or DEFAULT_SPAN_ORIGIN, + "status": status, + } # type: dict[str, Any] + + if span.attributes: + trace_context["data"] = dict(span.attributes) + + trace_context["dynamic_sampling_context"] = dsc_from_trace_state(span.context.trace_state) + + # TODO-neel-potel profiler thread_id, thread_name + + return trace_context + + +def trace_state_from_baggage(baggage): + # type: (Baggage) -> TraceState + items = [] + for k, v in baggage.sentry_items.items(): + key = Baggage.SENTRY_PREFIX + quote(k) + val = quote(str(v)) + items.append((key, val)) + return TraceState(items) + + +def serialize_trace_state(trace_state): + # type: (TraceState) -> str + sentry_items = [] + for k, v in trace_state.items(): + if Baggage.SENTRY_PREFIX_REGEX.match(k): + sentry_items.append((k, v)) + return ",".join(key + "=" + value for key, value in sentry_items) + + +def dsc_from_trace_state(trace_state): + # type: (TraceState) -> dict[str, str] + dsc = {} + for k, v in trace_state.items(): + if Baggage.SENTRY_PREFIX_REGEX.match(k): + key = re.sub(Baggage.SENTRY_PREFIX_REGEX, "", k) + dsc[key] = v + return dsc diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py index 775bdfaa7e..42f296a11d 100644 --- a/sentry_sdk/tracing.py +++ b/sentry_sdk/tracing.py @@ -1,7 +1,7 @@ import uuid import random import time -import warnings +import warnings from datetime import datetime, timedelta, timezone from opentelemetry import trace as otel_trace, context @@ -1447,7 +1447,14 @@ def start_child(self, **kwargs): def iter_headers(self): # type: () -> Iterator[Tuple[str, str]] - pass + yield SENTRY_TRACE_HEADER_NAME, self.to_traceparent() + + from sentry_sdk.integrations.opentelemetry.utils import ( + serialize_trace_state, + ) + + trace_state = self._otel_span.get_span_context().trace_state + yield BAGGAGE_HEADER_NAME, serialize_trace_state(trace_state) def to_traceparent(self): # type: () -> str @@ -1466,6 +1473,7 @@ def to_traceparent(self): def to_baggage(self): # type: () -> Optional[Baggage] + # TODO-neel-potel head SDK populate baggage mess pass def set_tag(self, key, value): @@ -1540,8 +1548,15 @@ def to_json(self): pass def get_trace_context(self): - # type: () -> Any - pass + # type: () -> dict[str, Any] + if not isinstance(self._otel_span, ReadableSpan): + return {} + + from sentry_sdk.integrations.opentelemetry.utils import ( + get_trace_context, + ) + + return get_trace_context(self._otel_span) def get_profile_context(self): # type: () -> Optional[ProfileContext]