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

This PR contains no functional changes. Its only purpose is to clearly
separate two enrichment steps in the segment consumer:

1. Enrich spans with information from the span tree (inherit from
segment, exclusive time from children, etc.)
2. Shim the Span V1 format to be compatible with the events processing
pipeline, for logic that is shared between the two pipelines.

Author: jjbayer

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 TreeSpan, 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) -> TreeSpan:
         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[TreeSpan]]:
         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 | TreeSpan) -> 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"],