ref(spans): Separate span tree enrichment from shimming (#98786)

Same as #98693 (which was
reverted), but with a fix to prevent invalid `dict` creation.

---------

Co-authored-by: Sebastian Zivota <[email protected]>

Author: 2 people

src/sentry/spans/consumers/process_segments/convert.py

@@ -8,7 +8,7 @@
 from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType
 from sentry_protos.snuba.v1.trace_item_pb2 import AnyValue, TraceItem
 
-from sentry.spans.consumers.process_segments.enrichment import Span
+from sentry.spans.consumers.process_segments.types import CompatibleSpan
 
 I64_MAX = 2**63 - 1
 
@@ -31,7 +31,7 @@
 }
 
 
-def convert_span_to_item(span: Span) -> TraceItem:
+def convert_span_to_item(span: CompatibleSpan) -> TraceItem:
     attributes: MutableMapping[str, AnyValue] = {}  # TODO
 
     client_sample_rate = 1.0

src/sentry/spans/consumers/process_segments/enrichment.py

@@ -1,9 +1,11 @@
 from collections import defaultdict
-from typing import Any, NotRequired
+from collections.abc import Sequence
+from typing import Any
 
 from sentry_kafka_schemas.schema_types.buffered_segments_v1 import SegmentSpan
 
 from sentry.performance_issues.types import SentryTags as PerformanceIssuesSentryTags
+from sentry.spans.consumers.process_segments.types import EnrichedSpan, get_span_op
 
 # Keys of shared sentry attributes that are shared across all spans in a segment. This list
 # is taken from `extract_shared_tags` in Relay.
@@ -43,27 +45,6 @@
 DEFAULT_SPAN_OP = "default"
 
 
-class Span(SegmentSpan, total=True):
-    """
-    Enriched version of the incoming span payload that has additional attributes
-    extracted.
-    """
-
-    # Added in enrichment
-    exclusive_time: float
-    exclusive_time_ms: float
-    op: str
-
-    sentry_tags: dict[str, Any]  # type: ignore[misc]  # XXX: fix w/ TypedDict extra_items once available
-
-    # Added by `SpanGroupingResults.write_to_spans` in `_enrich_spans`
-    hash: NotRequired[str]
-
-
-def _get_span_op(span: SegmentSpan | Span) -> str:
-    return span.get("data", {}).get("sentry.op") or DEFAULT_SPAN_OP
-
-
 def _find_segment_span(spans: list[SegmentSpan]) -> SegmentSpan | None:
     """
     Finds the segment in the span in the list that has ``is_segment`` set to
@@ -83,7 +64,9 @@ def _find_segment_span(spans: list[SegmentSpan]) -> SegmentSpan | None:
     return None
 
 
-class Enricher:
+class TreeEnricher:
+    """Enriches spans with information from their parent, child and sibling spans."""
+
     def __init__(self, spans: list[SegmentSpan]) -> None:
         self._segment_span = _find_segment_span(spans)
 
@@ -172,41 +155,28 @@ def _exclusive_time(self, span: SegmentSpan) -> float:
 
         return exclusive_time_us / 1_000
 
-    def enrich_span(self, span: SegmentSpan) -> Span:
+    def enrich_span(self, span: SegmentSpan) -> EnrichedSpan:
         exclusive_time = self._exclusive_time(span)
         data = self._data(span)
-        sentry_tags = self._sentry_tags(data)
         return {
             **span,
-            # Creates attributes for EAP spans that are required by logic shared with the
-            # event pipeline.
-            #
-            # Spans in the transaction event protocol had a slightly different schema
-            # compared to raw spans on the EAP topic. This function adds the missing
-            # attributes to the spans to make them compatible with the event pipeline
-            # logic.
             "data": data,
-            "sentry_tags": sentry_tags,
-            "op": _get_span_op(span),
-            # Note: Event protocol spans expect `exclusive_time` while EAP expects
-            # `exclusive_time_ms`. Both are the same value in milliseconds
-            "exclusive_time": exclusive_time,
             "exclusive_time_ms": exclusive_time,
         }
 
     @classmethod
-    def enrich_spans(cls, spans: list[SegmentSpan]) -> tuple[Span | None, list[Span]]:
+    def enrich_spans(cls, spans: list[SegmentSpan]) -> tuple[int | None, list[EnrichedSpan]]:
         inst = cls(spans)
         ret = []
-        segment_span = None
+        segment_idx = None
 
-        for span in spans:
+        for i, span in enumerate(spans):
             enriched = inst.enrich_span(span)
             if span is inst._segment_span:
-                segment_span = enriched
+                segment_idx = i
             ret.append(enriched)
 
-        return segment_span, ret
+        return segment_idx, ret
 
 
 def _get_mobile_start_type(segment: SegmentSpan) -> str | None:
@@ -226,12 +196,12 @@ def _get_mobile_start_type(segment: SegmentSpan) -> str | None:
 
 def _timestamp_by_op(spans: list[SegmentSpan], op: str) -> float | None:
     for span in spans:
-        if _get_span_op(span) == op:
+        if get_span_op(span) == op:
             return span["end_timestamp_precise"]
     return None
 
 
-def _span_interval(span: SegmentSpan | Span) -> tuple[int, int]:
+def _span_interval(span: SegmentSpan | EnrichedSpan) -> tuple[int, int]:
     """Get the start and end timestamps of a span in microseconds."""
     return _us(span["start_timestamp_precise"]), _us(span["end_timestamp_precise"])
 
@@ -243,7 +213,7 @@ def _us(timestamp: float) -> int:
 
 
 def compute_breakdowns(
-    spans: list[Span],
+    spans: Sequence[SegmentSpan],
     breakdowns_config: dict[str, dict[str, Any]],
 ) -> dict[str, float]:
     """
@@ -269,14 +239,14 @@ def compute_breakdowns(
     return ret
 
 
-def _compute_span_ops(spans: list[Span], config: Any) -> dict[str, float]:
+def _compute_span_ops(spans: Sequence[SegmentSpan], config: Any) -> dict[str, float]:
     matches = config.get("matches")
     if not matches:
         return {}
 
     intervals_by_op = defaultdict(list)
     for span in spans:
-        op = _get_span_op(span)
+        op = get_span_op(span)
         if operation_name := next(filter(lambda m: op.startswith(m), matches), None):
             intervals_by_op[operation_name].append(_span_interval(span))
 

src/sentry/spans/consumers/process_segments/factory.py

@@ -17,8 +17,8 @@
 from sentry import options
 from sentry.conf.types.kafka_definition import Topic
 from sentry.spans.consumers.process_segments.convert import convert_span_to_item
-from sentry.spans.consumers.process_segments.enrichment import Span
 from sentry.spans.consumers.process_segments.message import process_segment
+from sentry.spans.consumers.process_segments.types import CompatibleSpan
 from sentry.utils.arroyo import MultiprocessingPool, run_task_with_multiprocessing
 from sentry.utils.arroyo_producer import get_arroyo_producer
 from sentry.utils.kafka_config import get_kafka_producer_cluster_options, get_topic_definition
@@ -135,7 +135,7 @@ def _process_message(
         raise InvalidMessage(message.value.partition, message.value.offset)
 
 
-def _serialize_payload(span: Span, timestamp: datetime | None) -> Value[KafkaPayload]:
+def _serialize_payload(span: CompatibleSpan, timestamp: datetime | None) -> Value[KafkaPayload]:
     item = convert_span_to_item(span)
     return Value(
         KafkaPayload(

src/sentry/spans/consumers/process_segments/message.py

@@ -1,8 +1,8 @@
 import logging
 import types
 import uuid
-from copy import deepcopy
-from typing import Any, cast
+from collections.abc import Sequence
+from typing import cast
 
 from django.core.exceptions import ValidationError
 from sentry_kafka_schemas.schema_types.buffered_segments_v1 import SegmentSpan
@@ -26,7 +26,9 @@
 from sentry.receivers.features import record_generic_event_processed
 from sentry.receivers.onboarding import record_release_received
 from sentry.signals import first_insight_span_received, first_transaction_received
-from sentry.spans.consumers.process_segments.enrichment import Enricher, Span, compute_breakdowns
+from sentry.spans.consumers.process_segments.enrichment import TreeEnricher, compute_breakdowns
+from sentry.spans.consumers.process_segments.shim import build_shim_event_data, make_compatible
+from sentry.spans.consumers.process_segments.types import CompatibleSpan
 from sentry.spans.grouping.api import load_span_grouping_config
 from sentry.utils import metrics
 from sentry.utils.dates import to_datetime
@@ -39,7 +41,9 @@
 
 
 @metrics.wraps("spans.consumers.process_segments.process_segment")
-def process_segment(unprocessed_spans: list[SegmentSpan], skip_produce: bool = False) -> list[Span]:
+def process_segment(
+    unprocessed_spans: list[SegmentSpan], skip_produce: bool = False
+) -> list[CompatibleSpan]:
     segment_span, spans = _enrich_spans(unprocessed_spans)
     if segment_span is None:
         return spans
@@ -69,7 +73,9 @@ def process_segment(unprocessed_spans: list[SegmentSpan], skip_produce: bool = F
 
 
 @metrics.wraps("spans.consumers.process_segments.enrich_spans")
-def _enrich_spans(unprocessed_spans: list[SegmentSpan]) -> tuple[Span | None, list[Span]]:
+def _enrich_spans(
+    unprocessed_spans: list[SegmentSpan],
+) -> tuple[CompatibleSpan | None, list[CompatibleSpan]]:
     """
     Enriches all spans with data derived from the span tree and the segment.
 
@@ -80,7 +86,11 @@ def _enrich_spans(unprocessed_spans: list[SegmentSpan]) -> tuple[Span | None, li
     Returns the segment span, if any, and the list of enriched spans.
     """
 
-    segment, spans = Enricher.enrich_spans(unprocessed_spans)
+    segment_idx, tree_spans = TreeEnricher.enrich_spans(unprocessed_spans)
+
+    # Set attributes that are needed by logic shared with the event processing pipeline.
+    spans = [make_compatible(span) for span in tree_spans]
+    segment = spans[segment_idx] if segment_idx is not None else None
 
     # Calculate grouping hashes for performance issue detection
     config = load_span_grouping_config()
@@ -91,14 +101,16 @@ def _enrich_spans(unprocessed_spans: list[SegmentSpan]) -> tuple[Span | None, li
 
 
 @metrics.wraps("spans.consumers.process_segments.compute_breakdowns")
-def _compute_breakdowns(segment: Span, spans: list[Span], project: Project) -> None:
+def _compute_breakdowns(
+    segment: CompatibleSpan, spans: Sequence[CompatibleSpan], project: Project
+) -> None:
     config = project.get_option("sentry:breakdowns")
     breakdowns = compute_breakdowns(spans, config)
     segment.setdefault("data", {}).update(breakdowns)
 
 
 @metrics.wraps("spans.consumers.process_segments.create_models")
-def _create_models(segment: Span, project: Project) -> None:
+def _create_models(segment: CompatibleSpan, project: Project) -> None:
     """
     Creates the Environment and Release models, along with the necessary
     relationships between them and the Project model.
@@ -144,11 +156,13 @@ def _create_models(segment: Span, project: Project) -> None:
 
 
 @metrics.wraps("spans.consumers.process_segments.detect_performance_problems")
-def _detect_performance_problems(segment_span: Span, spans: list[Span], project: Project) -> None:
+def _detect_performance_problems(
+    segment_span: CompatibleSpan, spans: list[CompatibleSpan], project: Project
+) -> None:
     if not options.get("spans.process-segments.detect-performance-problems.enable"):
         return
 
-    event_data = _build_shim_event_data(segment_span, spans)
+    event_data = build_shim_event_data(segment_span, spans)
     performance_problems = detect_performance_problems(event_data, project, standalone=True)
 
     if not segment_span.get("_performance_issues_spans"):
@@ -191,55 +205,10 @@ def _detect_performance_problems(segment_span: Span, spans: list[Span], project:
         )
 
 
-def _build_shim_event_data(segment_span: Span, spans: list[Span]) -> dict[str, Any]:
-    data = segment_span.get("data", {})
-
-    event: dict[str, Any] = {
-        "type": "transaction",
-        "level": "info",
-        "contexts": {
-            "trace": {
-                "trace_id": segment_span["trace_id"],
-                "type": "trace",
-                "op": data.get("sentry.transaction.op"),
-                "span_id": segment_span["span_id"],
-                "hash": segment_span["hash"],
-            },
-        },
-        "event_id": uuid.uuid4().hex,
-        "project_id": segment_span["project_id"],
-        "transaction": data.get("sentry.transaction"),
-        "release": data.get("sentry.release"),
-        "dist": data.get("sentry.dist"),
-        "environment": data.get("sentry.environment"),
-        "platform": data.get("sentry.platform"),
-        "tags": [["environment", data.get("sentry.environment")]],
-        "received": segment_span["received"],
-        "timestamp": segment_span["end_timestamp_precise"],
-        "start_timestamp": segment_span["start_timestamp_precise"],
-        "datetime": to_datetime(segment_span["end_timestamp_precise"]).strftime(
-            "%Y-%m-%dT%H:%M:%SZ"
-        ),
-        "spans": [],
-    }
-
-    if (profile_id := segment_span.get("profile_id")) is not None:
-        event["contexts"]["profile"] = {"profile_id": profile_id, "type": "profile"}
-
-    # Add legacy span attributes required only by issue detectors. As opposed to
-    # real event payloads, this also adds the segment span so detectors can run
-    # topological sorting on the span tree.
-    for span in spans:
-        event_span = cast(dict[str, Any], deepcopy(span))
-        event_span["start_timestamp"] = span["start_timestamp_precise"]
-        event_span["timestamp"] = span["end_timestamp_precise"]
-        event["spans"].append(event_span)
-
-    return event
-
-
 @metrics.wraps("spans.consumers.process_segments.record_signals")
-def _record_signals(segment_span: Span, spans: list[Span], project: Project) -> None:
+def _record_signals(
+    segment_span: CompatibleSpan, spans: list[CompatibleSpan], project: Project
+) -> None:
     data = segment_span.get("data", {})
 
     record_generic_event_processed(
@@ -271,7 +240,7 @@ def _record_signals(segment_span: Span, spans: list[Span], project: Project) ->
 
 
 @metrics.wraps("spans.consumers.process_segments.record_outcomes")
-def _track_outcomes(segment_span: Span, spans: list[Span]) -> None:
+def _track_outcomes(segment_span: CompatibleSpan, spans: list[CompatibleSpan]) -> None:
     if options.get("spans.process-segments.outcome-aggregator.enable"):
         outcome_aggregator.track_outcome_aggregated(
             org_id=segment_span["organization_id"],