feat(trace-eap-waterfall): Updating eap transaction re-parenting logic (#87529)

As we reparent the eap-transactions being rendered under a collapsed
eap-transaction node, during expansion, we only need to target the
direct eap-transaction children of that node.

Consider this trace data: 
`txn 1`
-- span 1
----- `txn 2`
------ span 2
-------- `txn 3`

Step 1: Initially the trace is loaded as
`txn 1`
-- `txn 2`
----- `txn 3`

Step 2: Expanding txn 1 should render

`txn 1`
-- span 1
---- `txn 2`
------ `txn 3`

Step 3: Collapsing txn 1 should render
`txn 1`
-- `txn 2`
----- `txn 3`

Without the changes in this PR, Step 3 would render the wrong hierarchy
below:
`txn 1`
-- `txn 2`
-- `txn 3`

---------

Co-authored-by: Abdullah Khan <[email protected]>

Author: 2 people

static/app/views/performance/newTraceDetails/traceModels/__snapshots__/traceTree.spec.tsx.snap

@@ -69,20 +69,31 @@ eap trace root
 "
 `;
 
-exports[`TraceTree eap trace correctly renders eap-transactions collapsed state 1`] = `
+exports[`TraceTree eap trace correctly renders eap-transactions toggle state 1`] = `
 "
 eap trace root
   span.op - span.description (eap-transaction)
     span.op - span.description (eap-transaction)
+      span.op - span.description (eap-transaction)
 "
 `;
 
-exports[`TraceTree eap trace correctly renders eap-transactions expanded state 1`] = `
+exports[`TraceTree eap trace correctly renders eap-transactions toggle state 2`] = `
 "
 eap trace root
   span.op - span.description (eap-transaction)
     span.op - span.description
       span.op - span.description (eap-transaction)
+        span.op - span.description (eap-transaction)
+"
+`;
+
+exports[`TraceTree eap trace correctly renders eap-transactions toggle state 3`] = `
+"
+eap trace root
+  span.op - span.description (eap-transaction)
+    span.op - span.description (eap-transaction)
+      span.op - span.description (eap-transaction)
 "
 `;
 

static/app/views/performance/newTraceDetails/traceModels/traceTree.spec.tsx

@@ -534,14 +534,14 @@ describe('TraceTree', () => {
       expect(findEAPSpanByEventId(tree, 'eap-span-2')?.expanded).toBe(true);
     });
 
-    it('correctly renders eap-transactions collapsed state', () => {
+    it('correctly renders eap-transactions toggle state', () => {
       const tree = TraceTree.FromTrace(
         makeEAPTrace([
           makeEAPSpan({
             event_id: 'eap-span-1',
             start_timestamp: start,
             end_timestamp: start + 2,
-            is_transaction: true,
+            is_transaction: true, // is a transaction
             parent_span_id: undefined,
             children: [
               makeEAPSpan({
@@ -555,9 +555,27 @@ describe('TraceTree', () => {
                     event_id: 'eap-span-3',
                     start_timestamp: start + 2,
                     end_timestamp: start + 3,
-                    is_transaction: true,
+                    is_transaction: true, // is a transaction
                     parent_span_id: 'eap-span-2',
-                    children: [],
+                    children: [
+                      makeEAPSpan({
+                        event_id: 'eap-span-4',
+                        start_timestamp: start + 3,
+                        end_timestamp: start + 4,
+                        is_transaction: false,
+                        parent_span_id: 'eap-span-3',
+                        children: [
+                          makeEAPSpan({
+                            event_id: 'eap-span-5',
+                            start_timestamp: start + 4,
+                            end_timestamp: start + 5,
+                            is_transaction: true, // is a transaction
+                            parent_span_id: 'eap-span-4',
+                            children: [],
+                          }),
+                        ],
+                      }),
+                    ],
                   }),
                 ],
               }),
@@ -566,45 +584,17 @@ describe('TraceTree', () => {
         ]),
         traceMetadata
       );
-      expect(tree.build().serialize()).toMatchSnapshot();
-    });
 
-    it('correctly renders eap-transactions expanded state', () => {
-      const tree = TraceTree.FromTrace(
-        makeEAPTrace([
-          makeEAPSpan({
-            event_id: 'eap-span-1',
-            start_timestamp: start,
-            end_timestamp: start + 2,
-            is_transaction: true,
-            parent_span_id: undefined,
-            children: [
-              makeEAPSpan({
-                event_id: 'eap-span-2',
-                start_timestamp: start + 1,
-                end_timestamp: start + 4,
-                is_transaction: false,
-                parent_span_id: 'eap-span-1',
-                children: [
-                  makeEAPSpan({
-                    event_id: 'eap-span-3',
-                    start_timestamp: start + 2,
-                    end_timestamp: start + 3,
-                    is_transaction: true,
-                    parent_span_id: 'eap-span-2',
-                    children: [],
-                  }),
-                ],
-              }),
-            ],
-          }),
-        ]),
-        traceMetadata
-      );
+      // Assert initial state
+      expect(tree.build().serialize()).toMatchSnapshot();
 
+      // Assert expaneded state
       const eapTxn = findEAPSpanByEventId(tree, 'eap-span-1');
       tree.expand(eapTxn!, true);
+      expect(tree.build().serialize()).toMatchSnapshot();
 
+      // Assert state upon collapsing
+      tree.expand(eapTxn!, false);
       expect(tree.build().serialize()).toMatchSnapshot();
     });
   });

static/app/views/performance/newTraceDetails/traceModels/traceTree.tsx

@@ -356,12 +356,13 @@ export class TraceTree extends TraceTreeEventDispatcher {
 
       parentNode.children.push(node);
 
+      // Since we are reparenting EAP transactions at this stage, we need to sort the children
+      if (isEAPTransactionNode(node)) {
+        parentNode.children.sort(traceChronologicalSort);
+      }
+
       if (node.value && 'children' in node.value) {
-        // EAP spans are not sorted by default
-        const children = node.value.children.sort(
-          (a, b) => a.start_timestamp - b.start_timestamp
-        );
-        for (const child of children) {
+        for (const child of node.value.children) {
           visit(node, child);
         }
       }
@@ -1380,15 +1381,15 @@ export class TraceTree extends TraceTreeEventDispatcher {
 
   static ReparentEAPTransactions(
     node: TraceTreeNode<TraceTree.EAPSpan>,
+    findEAPTransactions: (
+      n: TraceTreeNode<TraceTree.EAPSpan>
+    ) => Array<TraceTreeNode<TraceTree.EAPSpan>>,
     findNewParent: (
       t: TraceTreeNode<TraceTree.EAPSpan>
     ) => TraceTreeNode<TraceTree.NodeValue> | null
   ): void {
     // Find all embedded eap-transactions, excluding the node itself
-    const eapTransactions = TraceTree.FindAll(
-      node,
-      n => isEAPTransactionNode(n) && n !== node
-    );
+    const eapTransactions = findEAPTransactions(node);
 
     for (const t of eapTransactions) {
       if (isEAPTransactionNode(t)) {
@@ -1469,10 +1470,13 @@ export class TraceTree extends TraceTreeEventDispatcher {
       node.expanded = expanded;
 
       // When eap-transaction nodes are expanded, we need to reparent the transactions under
-      // the eap-spans (by their parent_span_id) that were previously hidden.
+      // the eap-spans (by their parent_span_id) that were previously hidden. Note that this only impacts the
+      // direct eap-transaction children of the targetted eap-transaction node.
       if (isEAPTransactionNode(node)) {
-        TraceTree.ReparentEAPTransactions(node, t =>
-          TraceTree.FindByID(node, t.value.parent_span_id)
+        TraceTree.ReparentEAPTransactions(
+          node,
+          t => t.children.filter(c => isEAPTransactionNode(c)),
+          t => TraceTree.FindByID(node, t.value.parent_span_id)
         );
       }
 
@@ -1489,8 +1493,13 @@ export class TraceTree extends TraceTreeEventDispatcher {
       // Reparent the transactions from under the eap-spans in the expanded state, to under the closest eap-transaction
       // in the collapsed state.
       if (isEAPTransactionNode(node)) {
-        TraceTree.ReparentEAPTransactions(node, t =>
-          TraceTree.ParentEAPTransaction(t.parent)
+        TraceTree.ReparentEAPTransactions(
+          node,
+          t =>
+            TraceTree.FindAll(t, n => isEAPTransactionNode(n) && n !== t) as Array<
+              TraceTreeNode<TraceTree.EAPSpan>
+            >,
+          t => TraceTree.ParentEAPTransaction(t)
         );
       }