From 41d394c59434f4902d94cd4bf90232779d804925 Mon Sep 17 00:00:00 2001
From: Andre Pereira <adrapereira@gmail.com>
Date: Mon, 30 Oct 2023 14:26:33 +0000
Subject: [PATCH] Trace View: Fix issue with critical path performance (#77373)

Store child spans in Map for faster access on critical path render
---
 .../VirtualizedTraceView.tsx                  | 26 ++++++++++++++++---
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/public/app/features/explore/TraceView/components/TraceTimelineViewer/VirtualizedTraceView.tsx b/public/app/features/explore/TraceView/components/TraceTimelineViewer/VirtualizedTraceView.tsx
index 1d04158084b..7c45e88f548 100644
--- a/public/app/features/explore/TraceView/components/TraceTimelineViewer/VirtualizedTraceView.tsx
+++ b/public/app/features/explore/TraceView/components/TraceTimelineViewer/VirtualizedTraceView.tsx
@@ -211,9 +211,23 @@ function generateRowStatesFromTrace(
     : [];
 }
 
+function childSpansMap(trace: Trace | TNil) {
+  const childSpansMap = new Map<string, string[]>();
+  if (!trace) {
+    return childSpansMap;
+  }
+  trace.spans.forEach((span) => {
+    if (span.childSpanIds.length) {
+      childSpansMap.set(span.spanID, span.childSpanIds);
+    }
+  });
+  return childSpansMap;
+}
+
 const memoizedGenerateRowStates = memoizeOne(generateRowStatesFromTrace);
 const memoizedViewBoundsFunc = memoizeOne(createViewedBoundsFunc, isEqual);
 const memoizedGetClipping = memoizeOne(getClipping, isEqual);
+const memoizedChildSpansMap = memoizeOne(childSpansMap);
 
 // export from tests
 export class UnthemedVirtualizedTraceView extends React.Component<VirtualizedTraceViewProps> {
@@ -290,6 +304,10 @@ export class UnthemedVirtualizedTraceView extends React.Component<VirtualizedTra
     });
   }
 
+  getChildSpansMap() {
+    return memoizedChildSpansMap(this.props.trace);
+  }
+
   getAccessors() {
     const lv = this.listView;
     if (!lv) {
@@ -461,10 +479,10 @@ export class UnthemedVirtualizedTraceView extends React.Component<VirtualizedTra
     // This function called recursively to find all descendants of a span
     const findAllDescendants = (currentChildSpanIds: string[]) => {
       currentChildSpanIds.forEach((eachId) => {
-        const currentChildSpan = trace.spans.find((a) => a.spanID === eachId)!;
-        if (currentChildSpan.hasChildren) {
-          allChildSpanIds.push(...currentChildSpan.childSpanIds);
-          findAllDescendants(currentChildSpan.childSpanIds);
+        const childrenOfCurrent = this.getChildSpansMap().get(eachId);
+        if (childrenOfCurrent?.length) {
+          allChildSpanIds.push(...childrenOfCurrent);
+          findAllDescendants(childrenOfCurrent);
         }
       });
     };