fix(trace-view): Support multiple roots for the trace view (#25734)

* fix(trace-view): Support multiple roots for the trace view

- Currently the trace endpoint only returns the first root, this adds
  support for traces with multiple roots

bin/mock-traces

@@ -158,6 +158,29 @@ def main(slow=False):
             ],
         },
     )
+    print(f"    > Loading trace with many roots")  # NOQA
+    trace_id = uuid4().hex
+    for _ in range(15):
+        create_trace(
+            slow,
+            timestamp - timedelta(milliseconds=random_normal(4000, 250, 1000)),
+            timestamp,
+            generate_user(),
+            trace_id,
+            None,
+            {
+                "project": frontend_project,
+                "transaction": "/multiple-root/:root/",
+                "frontend": True,
+                "children": [
+                    {
+                        "project": backend_project,
+                        "transaction": "/multiple-root/child/",
+                        "children": [],
+                    }
+                ],
+            },
+        )
 
     print(f"    > Loading chained trace with orphans")  # NOQA
     trace_id = uuid4().hex

src/sentry/api/endpoints/organization_events_trace.py

@@ -355,24 +355,22 @@ class OrganizationEventsTraceEndpointBase(OrganizationEventsV2EndpointBase):  #
 
         warning_extra: Dict[str, str] = {"trace": trace_id, "organization": organization}
 
-        root = transactions[0] if is_root(transactions[0]) else None
-
-        # Look for extra roots
-        extra_roots = 0
-        for item in transactions[1:]:
+        # Look for the roots
+        roots: List[SnubaTransaction] = []
+        for item in transactions:
             if is_root(item):
-                extra_roots += 1
+                roots.append(item)
             else:
                 break
-        if extra_roots > 0:
+        if len(roots) > 1:
             sentry_sdk.set_tag("discover.trace-view.warning", "root.extra-found")
             logger.warning(
                 "discover.trace-view.root.extra-found",
-                {"extra_roots": extra_roots, **warning_extra},
+                {"extra_roots": len(roots), **warning_extra},
             )
 
         return Response(
-            self.serialize(transactions, errors, root, warning_extra, event_id, detailed)
+            self.serialize(transactions, errors, roots, warning_extra, event_id, detailed)
         )
 
 
@@ -430,7 +428,7 @@ class OrganizationEventsTraceLightEndpoint(OrganizationEventsTraceEndpointBase):
         self,
         transactions: Sequence[SnubaTransaction],
         errors: Sequence[SnubaError],
-        root: Optional[SnubaTransaction],
+        roots: Sequence[SnubaTransaction],
         warning_extra: Dict[str, str],
         event_id: Optional[str],
         detailed: bool = False,
@@ -446,31 +444,38 @@ class OrganizationEventsTraceLightEndpoint(OrganizationEventsTraceEndpointBase):
         root_id: Optional[str] = None
 
         with sentry_sdk.start_span(op="building.trace", description="light trace"):
-            # We might not be necessarily connected to the root if we're on an orphan event
-            if root is not None and root["id"] != snuba_event["id"]:
-                # Get the root event and see if the current event's span is in the root event
-                root_event = eventstore.get_event_by_id(root["project.id"], root["id"])
-                root_spans: NodeSpans = root_event.data.get("spans", [])
-                root_span = find_event(
-                    root_spans,
-                    lambda item: item is not None
-                    and item["span_id"] == snuba_event["trace.parent_span"],
-                )
+            # Going to nodestore is more expensive than looping twice so check if we're on the root first
+            for root in roots:
+                if root["id"] == snuba_event["id"]:
+                    current_generation = 0
+                    break
 
-                # We only know to add the root if its the direct parent
-                if root_span is not None:
-                    # For the light response, the parent will be unknown unless it is a direct descendent of the root
-                    root_id = root["id"]
-                    trace_results.append(
-                        TraceEvent(
-                            root,
-                            None,
-                            0,
+            if current_generation is None:
+                for root in roots:
+                    # We might not be necessarily connected to the root if we're on an orphan event
+                    if root["id"] != snuba_event["id"]:
+                        # Get the root event and see if the current event's span is in the root event
+                        root_event = eventstore.get_event_by_id(root["project.id"], root["id"])
+                        root_spans: NodeSpans = root_event.data.get("spans", [])
+                        root_span = find_event(
+                            root_spans,
+                            lambda item: item is not None
+                            and item["span_id"] == snuba_event["trace.parent_span"],
                         )
-                    )
-                    current_generation = 1
-            elif root is not None and root["id"] == snuba_event["id"]:
-                current_generation = 0
+
+                        # We only know to add the root if its the direct parent
+                        if root_span is not None:
+                            # For the light response, the parent will be unknown unless it is a direct descendent of the root
+                            root_id = root["id"]
+                            trace_results.append(
+                                TraceEvent(
+                                    root,
+                                    None,
+                                    0,
+                                )
+                            )
+                            current_generation = 1
+                            break
 
             current_event = TraceEvent(snuba_event, root_id, current_generation)
             trace_results.append(current_event)
@@ -527,7 +532,7 @@ class OrganizationEventsTraceEndpoint(OrganizationEventsTraceEndpointBase):
         self,
         transactions: Sequence[SnubaTransaction],
         errors: Sequence[SnubaError],
-        root: Optional[SnubaTransaction],
+        roots: Sequence[SnubaTransaction],
         warning_extra: Dict[str, str],
         event_id: Optional[str],
         detailed: bool = False,
@@ -545,9 +550,10 @@ class OrganizationEventsTraceEndpoint(OrganizationEventsTraceEndpointBase):
         # on python 3.7.
         results_map: Dict[Optional[str], List[TraceEvent]] = OrderedDict()
         to_check: Deque[SnubaTransaction] = deque()
-        if root:
+        results_map[None] = []
+        for root in roots:
             parent_events[root["id"]] = TraceEvent(root, None, 0)
-            results_map[None] = [parent_events[root["id"]]]
+            results_map[None].append(parent_events[root["id"]])
             to_check.append(root)
 
         with sentry_sdk.start_span(op="building.trace", description="full trace"):

tests/snuba/api/endpoints/test_organization_events_trace.py

@@ -296,30 +296,48 @@ class OrganizationEventsTraceLightEndpointTest(OrganizationEventsTraceEndpointBa
         assert event["parent_span_id"] == parent_span_id
         assert event["event_id"] == no_root_event.event_id
 
-    def test_multiple_roots(self):
-        self.create_event(
-            trace=self.trace_id,
-            transaction="/second_root",
-            spans=[],
-            parent_span_id=None,
-            project_id=self.project.id,
-        )
+    def test_root_event(self):
+        root_event_id = self.root_event.event_id
+
         with self.feature(self.FEATURES):
             response = self.client.get(
                 self.url,
-                data={"event_id": self.root_event.event_id},
+                data={"event_id": root_event_id, "project": -1},
                 format="json",
             )
 
         assert response.status_code == 200, response.content
 
-    def test_root_event(self):
-        root_event_id = self.root_event.event_id
+        assert len(response.data) == 4
+        events = {item["event_id"]: item for item in response.data}
+
+        assert root_event_id in events
+        event = events[root_event_id]
+        assert event["generation"] == 0
+        assert event["parent_event_id"] is None
+        assert event["parent_span_id"] is None
+
+        for i, child_event in enumerate(self.gen1_events):
+            child_event_id = child_event.event_id
+            assert child_event_id in events
+            event = events[child_event_id]
+            assert event["generation"] == 1
+            assert event["parent_event_id"] == root_event_id
+            assert event["parent_span_id"] == self.root_span_ids[i]
 
+    def test_root_with_multiple_roots(self):
+        root_event_id = self.root_event.event_id
+        self.create_event(
+            trace=self.trace_id,
+            transaction="/second_root",
+            spans=[],
+            parent_span_id=None,
+            project_id=self.project.id,
+        )
         with self.feature(self.FEATURES):
             response = self.client.get(
                 self.url,
-                data={"event_id": root_event_id, "project": -1},
+                data={"event_id": self.root_event.event_id},
                 format="json",
             )
 
@@ -377,6 +395,48 @@ class OrganizationEventsTraceLightEndpointTest(OrganizationEventsTraceEndpointBa
         assert event["parent_event_id"] == current_event
         assert event["parent_span_id"] == self.gen1_span_ids[0]
 
+    def test_direct_parent_with_children_and_multiple_root(self):
+        root_event_id = self.root_event.event_id
+        current_event = self.gen1_events[0].event_id
+        child_event_id = self.gen2_events[0].event_id
+        self.create_event(
+            trace=self.trace_id,
+            transaction="/second_root",
+            spans=[],
+            parent_span_id=None,
+            project_id=self.project.id,
+        )
+
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={"event_id": current_event, "project": -1},
+                format="json",
+            )
+
+        assert response.status_code == 200, response.content
+
+        assert len(response.data) == 3
+        events = {item["event_id"]: item for item in response.data}
+
+        assert root_event_id in events
+        event = events[root_event_id]
+        assert event["generation"] == 0
+        assert event["parent_event_id"] is None
+        assert event["parent_span_id"] is None
+
+        assert current_event in events
+        event = events[current_event]
+        assert event["generation"] == 1
+        assert event["parent_event_id"] == root_event_id
+        assert event["parent_span_id"] == self.root_span_ids[0]
+
+        assert child_event_id in events
+        event = events[child_event_id]
+        assert event["generation"] == 2
+        assert event["parent_event_id"] == current_event
+        assert event["parent_span_id"] == self.gen1_span_ids[0]
+
     def test_second_generation_with_children(self):
         current_event = self.gen2_events[0].event_id
         child_event_id = self.gen3_event.event_id
@@ -698,6 +758,40 @@ class OrganizationEventsTraceEndpointTest(OrganizationEventsTraceEndpointBase):
         # We didn't even try to start the loop of spans
         assert len(gen3_1["children"]) == 0
 
+    def test_multiple_roots(self):
+        trace_id = uuid4().hex
+        first_root = self.create_event(
+            trace=trace_id,
+            transaction="/first_root",
+            spans=[],
+            parent_span_id=None,
+            project_id=self.project.id,
+            duration=1000,
+        )
+        second_root = self.create_event(
+            trace=trace_id,
+            transaction="/second_root",
+            spans=[],
+            parent_span_id=None,
+            project_id=self.project.id,
+            duration=500,
+        )
+        self.url = reverse(
+            self.url_name,
+            kwargs={"organization_slug": self.project.organization.slug, "trace_id": trace_id},
+        )
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={"project": -1},
+                format="json",
+            )
+
+        assert response.status_code == 200, response.content
+        assert len(response.data) == 2
+        self.assert_event(response.data[0], first_root, "first_root")
+        self.assert_event(response.data[1], second_root, "second_root")
+
     def test_sibling_transactions(self):
         """ More than one transaction can share a parent_span_id """
         gen3_event_siblings = [