[DevTools][Timeline Profiler] Component Stacks Backend (#24776)

This PR adds a component stack field to the `schedule-state-update` event. The algorithm is as follows:
* During profiling, whenever a state update happens collect the parents of the fiber that caused the state update and store it in a map
* After profiling finishes, post process the `schedule-state-update` event and using the parent fibers, generate the component stack by using`describeFiber`, a function that uses error throwing to get the location of the component by calling the component without props.

---

Co-authored-by: Blake Friedman <blake.friedman@gmail.com>

packages/react-devtools-shared/src/__tests__/TimelineProfiler-test.js

@@ -9,6 +9,15 @@
 
 'use strict';
 
+function normalizeCodeLocInfo(str) {
+  return (
+    typeof str === 'string' &&
+    str.replace(/\n +(?:at|in) ([\S]+)[^\n]*/g, function(m, name) {
+      return '\n    in ' + name + ' (at **)';
+    })
+  );
+}
+
 describe('Timeline profiler', () => {
   let React;
   let ReactDOMClient;
@@ -1175,6 +1184,18 @@ describe('Timeline profiler', () => {
         if (timelineData) {
           expect(timelineData).toHaveLength(1);
 
+          // normalize the location for component stack source
+          // for snapshot testing
+          timelineData.forEach(data => {
+            data.schedulingEvents.forEach(event => {
+              if (event.componentStack) {
+                event.componentStack = normalizeCodeLocInfo(
+                  event.componentStack,
+                );
+              }
+            });
+          });
+
           return timelineData[0];
         } else {
           return null;
@@ -1256,6 +1277,8 @@ describe('Timeline profiler', () => {
           Array [
             Object {
               "componentName": "Example",
+              "componentStack": "
+              in Example (at **)",
               "lanes": "0b0000000000000000000000000000100",
               "timestamp": 10,
               "type": "schedule-state-update",
@@ -1263,6 +1286,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "Example",
+              "componentStack": "
+              in Example (at **)",
               "lanes": "0b0000000000000000000000001000000",
               "timestamp": 10,
               "type": "schedule-state-update",
@@ -1270,6 +1295,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "Example",
+              "componentStack": "
+              in Example (at **)",
               "lanes": "0b0000000000000000000000001000000",
               "timestamp": 10,
               "type": "schedule-state-update",
@@ -1277,6 +1304,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "Example",
+              "componentStack": "
+              in Example (at **)",
               "lanes": "0b0000000000000000000000000010000",
               "timestamp": 10,
               "type": "schedule-state-update",
@@ -1614,6 +1643,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "Example",
+              "componentStack": "
+              in Example (at **)",
               "lanes": "0b0000000000000000000000000000001",
               "timestamp": 20,
               "type": "schedule-state-update",
@@ -1741,6 +1772,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "Example",
+              "componentStack": "
+              in Example (at **)",
               "lanes": "0b0000000000000000000000000010000",
               "timestamp": 10,
               "type": "schedule-state-update",
@@ -1872,6 +1905,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "Example",
+              "componentStack": "
+              in Example (at **)",
               "lanes": "0b0000000000000000000000000000001",
               "timestamp": 21,
               "type": "schedule-state-update",
@@ -1934,6 +1969,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "Example",
+              "componentStack": "
+              in Example (at **)",
               "lanes": "0b0000000000000000000000000010000",
               "timestamp": 21,
               "type": "schedule-state-update",
@@ -1982,6 +2019,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "Example",
+              "componentStack": "
+              in Example (at **)",
               "lanes": "0b0000000000000000000000000010000",
               "timestamp": 20,
               "type": "schedule-state-update",
@@ -2065,6 +2104,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "ErrorBoundary",
+              "componentStack": "
+              in ErrorBoundary (at **)",
               "lanes": "0b0000000000000000000000000000001",
               "timestamp": 20,
               "type": "schedule-state-update",
@@ -2177,6 +2218,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "ErrorBoundary",
+              "componentStack": "
+              in ErrorBoundary (at **)",
               "lanes": "0b0000000000000000000000000000001",
               "timestamp": 30,
               "type": "schedule-state-update",
@@ -2441,6 +2484,52 @@ describe('Timeline profiler', () => {
           }
         `);
       });
+
+      it('should generate component stacks for state update', async () => {
+        function CommponentWithChildren({initialRender}) {
+          Scheduler.unstable_yieldValue('Render ComponentWithChildren');
+          return <Child initialRender={initialRender} />;
+        }
+
+        function Child({initialRender}) {
+          const [didRender, setDidRender] = React.useState(initialRender);
+          if (!didRender) {
+            setDidRender(true);
+          }
+          Scheduler.unstable_yieldValue('Render Child');
+          return null;
+        }
+
+        renderRootHelper(<CommponentWithChildren initialRender={false} />);
+
+        expect(Scheduler).toFlushAndYield([
+          'Render ComponentWithChildren',
+          'Render Child',
+          'Render Child',
+        ]);
+
+        const timelineData = stopProfilingAndGetTimelineData();
+        expect(timelineData.schedulingEvents).toMatchInlineSnapshot(`
+          Array [
+            Object {
+              "lanes": "0b0000000000000000000000000010000",
+              "timestamp": 10,
+              "type": "schedule-render",
+              "warning": null,
+            },
+            Object {
+              "componentName": "Child",
+              "componentStack": "
+              in Child (at **)
+              in CommponentWithChildren (at **)",
+              "lanes": "0b0000000000000000000000000010000",
+              "timestamp": 10,
+              "type": "schedule-state-update",
+              "warning": null,
+            },
+          ]
+        `);
+      });
     });
 
     describe('when not profiling', () => {

packages/react-devtools-shared/src/__tests__/preprocessData-test.js

@@ -9,6 +9,15 @@
 
 'use strict';
 
+function normalizeCodeLocInfo(str) {
+  return (
+    typeof str === 'string' &&
+    str.replace(/\n +(?:at|in) ([\S]+)[^\n]*/g, function(m, name) {
+      return '\n    in ' + name + ' (at **)';
+    })
+  );
+}
+
 describe('Timeline profiler', () => {
   let React;
   let ReactDOM;
@@ -2134,6 +2143,15 @@ describe('Timeline profiler', () => {
       const data = store.profilerStore.profilingData?.timelineData;
       expect(data).toHaveLength(1);
       const timelineData = data[0];
+
+      // normalize the location for component stack source
+      // for snapshot testing
+      timelineData.schedulingEvents.forEach(event => {
+        if (event.componentStack) {
+          event.componentStack = normalizeCodeLocInfo(event.componentStack);
+        }
+      });
+
       expect(timelineData).toMatchInlineSnapshot(`
         Object {
           "batchUIDToMeasuresMap": Map {
@@ -2415,6 +2433,8 @@ describe('Timeline profiler', () => {
             },
             Object {
               "componentName": "App",
+              "componentStack": "
+            in App (at **)",
               "lanes": "0b0000000000000000000000000010000",
               "timestamp": 10,
               "type": "schedule-state-update",

packages/react-devtools-shared/src/backend/DevToolsFiberComponentStack.js

@@ -21,7 +21,7 @@ import {
   describeClassComponentFrame,
 } from './DevToolsComponentStackFrame';
 
-function describeFiber(
+export function describeFiber(
   workTagMap: WorkTagMap,
   workInProgress: Fiber,
   currentDispatcherRef: CurrentDispatcherRef,

packages/react-devtools-shared/src/backend/profilingHooks.js

@@ -11,6 +11,8 @@ import type {
   Lane,
   Lanes,
   DevToolsProfilingHooks,
+  WorkTagMap,
+  CurrentDispatcherRef,
 } from 'react-devtools-shared/src/backend/types';
 import type {Fiber} from 'react-reconciler/src/ReactInternalTypes';
 import type {Wakeable} from 'shared/ReactTypes';
@@ -22,6 +24,8 @@ import type {
   ReactMeasureType,
   TimelineData,
   SuspenseEvent,
+  SchedulingEvent,
+  ReactScheduleStateUpdateEvent,
 } from 'react-devtools-timeline/src/types';
 
 import isArray from 'shared/isArray';
@@ -29,6 +33,7 @@ import {
   REACT_TOTAL_NUM_LANES,
   SCHEDULING_PROFILER_VERSION,
 } from 'react-devtools-timeline/src/constants';
+import {describeFiber} from './DevToolsFiberComponentStack';
 
 // Add padding to the start/stop time of the profile.
 // This makes the UI nicer to use.
@@ -98,17 +103,22 @@ export function createProfilingHooks({
   getDisplayNameForFiber,
   getIsProfiling,
   getLaneLabelMap,
+  workTagMap,
+  currentDispatcherRef,
   reactVersion,
 }: {|
   getDisplayNameForFiber: (fiber: Fiber) => string | null,
   getIsProfiling: () => boolean,
   getLaneLabelMap?: () => Map<Lane, string> | null,
+  currentDispatcherRef?: CurrentDispatcherRef,
+  workTagMap: WorkTagMap,
   reactVersion: string,
 |}): Response {
   let currentBatchUID: BatchUID = 0;
   let currentReactComponentMeasure: ReactComponentMeasure | null = null;
   let currentReactMeasuresStack: Array<ReactMeasure> = [];
   let currentTimelineData: TimelineData | null = null;
+  let currentFiberStacks: Map<SchedulingEvent, Array<Fiber>> = new Map();
   let isProfiling: boolean = false;
   let nextRenderShouldStartNewBatch: boolean = false;
 
@@ -774,6 +784,16 @@ export function createProfilingHooks({
     }
   }
 
+  function getParentFibers(fiber: Fiber): Array<Fiber> {
+    const parents = [];
+    let parent = fiber;
+    while (parent !== null) {
+      parents.push(parent);
+      parent = parent.return;
+    }
+    return parents;
+  }
+
   function markStateUpdateScheduled(fiber: Fiber, lane: Lane): void {
     if (isProfiling || supportsUserTimingV3) {
       const componentName = getDisplayNameForFiber(fiber) || 'Unknown';
@@ -781,13 +801,17 @@ export function createProfilingHooks({
       if (isProfiling) {
         // TODO (timeline) Record and cache component stack
         if (currentTimelineData) {
-          currentTimelineData.schedulingEvents.push({
+          const event: ReactScheduleStateUpdateEvent = {
             componentName,
+            // Store the parent fibers so we can post process
+            // them after we finish profiling
             lanes: laneToLanesArray(lane),
             timestamp: getRelativeTime(),
             type: 'schedule-state-update',
             warning: null,
-          });
+          };
+          currentFiberStacks.set(event, getParentFibers(fiber));
+          currentTimelineData.schedulingEvents.push(event);
         }
       }
 
@@ -831,6 +855,7 @@ export function createProfilingHooks({
         currentBatchUID = 0;
         currentReactComponentMeasure = null;
         currentReactMeasuresStack = [];
+        currentFiberStacks = new Map();
         currentTimelineData = {
           // Session wide metadata; only collected once.
           internalModuleSourceToRanges,
@@ -858,6 +883,30 @@ export function createProfilingHooks({
           snapshotHeight: 0,
         };
         nextRenderShouldStartNewBatch = true;
+      } else {
+        // Postprocess Profile data
+        if (currentTimelineData !== null) {
+          currentTimelineData.schedulingEvents.forEach(event => {
+            if (event.type === 'schedule-state-update') {
+              // TODO(luna): We can optimize this by creating a map of
+              // fiber to component stack instead of generating the stack
+              // for every fiber every time
+              const fiberStack = currentFiberStacks.get(event);
+              if (fiberStack && currentDispatcherRef != null) {
+                event.componentStack = fiberStack.reduce((trace, fiber) => {
+                  return (
+                    trace +
+                    describeFiber(workTagMap, fiber, currentDispatcherRef)
+                  );
+                }, '');
+              }
+            }
+          });
+        }
+
+        // Clear the current fiber stacks so we don't hold onto the fibers
+        // in memory after profiling finishes
+        currentFiberStacks.clear();
       }
     }
   }

packages/react-devtools-shared/src/backend/renderer.js

@@ -660,6 +660,8 @@ export function attach(
       getDisplayNameForFiber,
       getIsProfiling: () => isProfiling,
       getLaneLabelMap,
+      currentDispatcherRef: renderer.currentDispatcherRef,
+      workTagMap: ReactTypeOfWork,
       reactVersion: version,
     });
 

packages/react-devtools-timeline/src/types.js

@@ -51,6 +51,7 @@ export type ReactScheduleRenderEvent = {|
 |};
 export type ReactScheduleStateUpdateEvent = {|
   ...BaseReactScheduleEvent,
+  +componentStack?: string,
   +type: 'schedule-state-update',
 |};
 export type ReactScheduleForceUpdateEvent = {|

packages/shared/ReactComponentStackFrame.js

@@ -131,6 +131,9 @@ export function describeNativeComponentFrame(
       } catch (x) {
         control = x;
       }
+      // TODO(luna): This will currently only throw if the function component
+      // tries to access React/ReactDOM/props. We should probably make this throw
+      // in simple components too
       fn();
     }
   } catch (sample) {