feat(spans): Implement segment enrichment (#87273)

Ports the relevant bits from Relay's transaction-based normalization
logic into the segment consumer. Note that we intentionally keep all
normalization within spans in Relay. That is, we assume that Relay will
write normalized data attributes into spans and the segment consumer
merely has to move them to child spans and/or infer new attributes
across the span hierarchy.

To keep things simple, we pull normalized attributes from `sentry_tags`.
In a future refactor, this logic may be moved out from Relay, too, in
which case we will have to access `span.data` directly.

The final remaining function, `set_span_exclusive_time`, is going to be
added in a follow-up.

Author: jan-auer

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

@@ -32,6 +32,37 @@
 logger = logging.getLogger(__name__)
 
 
+# Keys in `sentry_tags` that are shared across all spans in a segment. This list
+# is taken from `extract_shared_tags` in Relay.
+SHARED_TAG_KEYS = (
+    "release",
+    "user",
+    "user.id",
+    "user.ip",
+    "user.username",
+    "user.email",
+    "user.geo.country_code",
+    "user.geo.subregion",
+    "environment",
+    "transaction",
+    "transaction.method",
+    "transaction.op",
+    "trace.status",
+    "mobile",
+    "os.name",
+    "device.class",
+    "browser.name",
+    "profiler_id",
+    "sdk.name",
+    "sdk.version",
+    "platform",
+    "thread.id",
+    "thread.name",
+)
+
+MOBILE_MAIN_THREAD_NAME = "main"
+
+
 class Span(SchemaSpan, total=False):
     start_timestamp_precise: float  # Missing in schema
     end_timestamp_precise: float  # Missing in schema
@@ -41,15 +72,14 @@ class Span(SchemaSpan, total=False):
 
 def process_segment(spans: list[Span]) -> list[Span]:
     segment_span = _find_segment_span(spans)
+    _enrich_spans(segment_span, spans)
+
     if segment_span is None:
-        # TODO: Handle segments without a defined segment span once all
-        # functions are refactored to a span interface.
         return spans
 
     with metrics.timer("spans.consumers.process_segments.get_project"):
         project = Project.objects.get_from_cache(id=segment_span["project_id"])
 
-    _enrich_spans(segment_span, spans)
     _create_models(segment_span, project)
     _detect_performance_problems(segment_span, spans, project)
     _record_signals(segment_span, spans, project)
@@ -85,16 +115,73 @@ def _find_segment_span(spans: list[Span]) -> Span | None:
 @metrics.wraps("spans.consumers.process_segments.enrich_spans")
 def _enrich_spans(segment: Span | None, spans: list[Span]) -> None:
     for span in spans:
-        span["op"] = span.get("sentry_tags", {}).get("op") or DEFAULT_SPAN_OP
+        # TODO: TEST THAT THIS RUNS WITHOUT A SEGMENT SPAN!
+        sentry_tags = span.setdefault("sentry_tags", {})
+        span["op"] = sentry_tags.get("op") or DEFAULT_SPAN_OP
+        # TODO: port set_span_exclusive_time
 
-        # TODO: Add Relay's enrichment here.
+    if segment:
+        _set_shared_tags(segment, spans)
 
     # Calculate grouping hashes for performance issue detection
     config = load_span_grouping_config()
     groupings = config.execute_strategy_standalone(spans)
     groupings.write_to_spans(spans)
 
 
+def _set_shared_tags(segment: Span, spans: list[Span]) -> None:
+    # Assume that Relay has extracted the shared tags into `sentry_tags` on the
+    # root span. Once `sentry_tags` is removed, the logic from
+    # `extract_shared_tags` should be moved here.
+    segment_tags = segment.get("sentry_tags", {})
+    shared_tags = {k: v for k, v in segment_tags.items() if k in SHARED_TAG_KEYS}
+
+    is_mobile = segment_tags.get("mobile") == "true"
+    mobile_start_type = _get_mobile_start_type(segment)
+    ttid_ts = _timestamp_by_op(spans, "ui.load.initial_display")
+    ttfd_ts = _timestamp_by_op(spans, "ui.load.full_display")
+
+    for span in spans:
+        span_tags = cast(dict[str, Any], span["sentry_tags"])
+
+        if is_mobile:
+            if span_tags.get("thread.name") == MOBILE_MAIN_THREAD_NAME:
+                span_tags["main_thread"] = "true"
+            if not span_tags.get("app_start_type") and mobile_start_type:
+                span_tags["app_start_type"] = mobile_start_type
+
+        if ttid_ts is not None and span["end_timestamp_precise"] <= ttid_ts:
+            span_tags["ttid"] = "ttid"
+        if ttfd_ts is not None and span["end_timestamp_precise"] <= ttfd_ts:
+            span_tags["ttfd"] = "ttfd"
+
+        for key, value in shared_tags.items():
+            if span_tags.get(key) is None:
+                span_tags[key] = value
+
+
+def _get_mobile_start_type(segment: Span) -> str | None:
+    """
+    Check the measurements on the span to determine what kind of start type the
+    event is.
+    """
+    measurements = segment.get("measurements") or {}
+
+    if "app_start_cold" in measurements:
+        return "cold"
+    if "app_start_warm" in measurements:
+        return "warm"
+
+    return None
+
+
+def _timestamp_by_op(spans: list[Span], op: str) -> float | None:
+    for span in spans:
+        if span["op"] == op:
+            return span["end_timestamp_precise"]
+    return None
+
+
 @metrics.wraps("spans.consumers.process_segments.create_models")
 def _create_models(segment: Span, project: Project) -> None:
     """

tests/sentry/spans/consumers/process_segments/test_message.py

@@ -29,6 +29,11 @@ def generate_basic_spans(self):
             start_timestamp_precise=1707953018.867,
         )
 
+        del child_span["sentry_tags"]["transaction"]
+        del child_span["sentry_tags"]["transaction.method"]
+        del child_span["sentry_tags"]["transaction.op"]
+        del child_span["sentry_tags"]["user"]
+
         return [child_span, segment_span]
 
     def generate_n_plus_one_spans(self):
@@ -71,6 +76,33 @@ def repeating_span():
 
         return spans
 
+    def test_enrich_spans(self):
+        spans = self.generate_basic_spans()
+        processed_spans = process_segment(spans)
+
+        assert len(processed_spans) == len(spans)
+        child_span, segment_span = processed_spans
+        child_tags = child_span["sentry_tags"]
+        segment_tags = segment_span["sentry_tags"]
+
+        assert child_tags["transaction"] == segment_tags["transaction"]
+        assert child_tags["transaction.method"] == segment_tags["transaction.method"]
+        assert child_tags["transaction.op"] == segment_tags["transaction.op"]
+        assert child_tags["user"] == segment_tags["user"]  # type: ignore[typeddict-item]
+
+    def test_enrich_spans_no_segment(self):
+        spans = self.generate_basic_spans()
+        for span in spans:
+            span["is_segment"] = False
+            del span["sentry_tags"]
+
+        processed_spans = process_segment(spans)
+        assert len(processed_spans) == len(spans)
+        for i, span in enumerate(processed_spans):
+            assert span["span_id"] == spans[i]["span_id"]
+            assert span["op"]
+            assert span["hash"]
+
     def test_create_models(self):
         spans = self.generate_basic_spans()
         assert process_segment(spans)
@@ -86,20 +118,6 @@ def test_create_models(self):
         )
         assert release.date_added.timestamp() == spans[0]["end_timestamp_precise"]
 
-    def test_empty_defaults(self):
-        spans = self.generate_basic_spans()
-        for span in spans:
-            del span["sentry_tags"]
-
-        processed_spans = process_segment(spans)
-        assert len(processed_spans) == len(spans)
-        assert processed_spans[0]["span_id"] == spans[0]["span_id"]
-        assert processed_spans[1]["span_id"] == spans[1]["span_id"]
-
-        # double-check that we actually ran through processing. The "op"
-        # attribute does not exist in the original spans.
-        assert processed_spans[0]["op"]
-
     @override_options(
         {
             "standalone-spans.detect-performance-problems.enable": True,