[Transition Tracing] Add Tracing Markers (#24686)

This PR adds support for Tracing Markers as well as onTracingMarkerComplete

packages/react-reconciler/src/ReactFiber.new.js

@@ -19,6 +19,7 @@ import type {
   OffscreenProps,
   OffscreenInstance,
 } from './ReactFiberOffscreenComponent';
+import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.new';
 
 import {
   createRootStrictEffectsByDefault,
@@ -757,6 +758,11 @@ export function createFiberFromTracingMarker(
   const fiber = createFiber(TracingMarkerComponent, pendingProps, key, mode);
   fiber.elementType = REACT_TRACING_MARKER_TYPE;
   fiber.lanes = lanes;
+  const tracingMarkerInstance: TracingMarkerInstance = {
+    transitions: null,
+    pendingSuspenseBoundaries: null,
+  };
+  fiber.stateNode = tracingMarkerInstance;
   return fiber;
 }
 

packages/react-reconciler/src/ReactFiber.old.js

@@ -19,6 +19,7 @@ import type {
   OffscreenProps,
   OffscreenInstance,
 } from './ReactFiberOffscreenComponent';
+import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.old';
 
 import {
   createRootStrictEffectsByDefault,
@@ -757,6 +758,11 @@ export function createFiberFromTracingMarker(
   const fiber = createFiber(TracingMarkerComponent, pendingProps, key, mode);
   fiber.elementType = REACT_TRACING_MARKER_TYPE;
   fiber.lanes = lanes;
+  const tracingMarkerInstance: TracingMarkerInstance = {
+    transitions: null,
+    pendingSuspenseBoundaries: null,
+  };
+  fiber.stateNode = tracingMarkerInstance;
   return fiber;
 }
 

packages/react-reconciler/src/ReactFiberBeginWork.new.js

@@ -35,6 +35,7 @@ import type {
 } from './ReactFiberCacheComponent.new';
 import type {UpdateQueue} from './ReactFiberClassUpdateQueue.new';
 import type {RootState} from './ReactFiberRoot.new';
+import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.new';
 import {
   enableSuspenseAvoidThisFallback,
   enableCPUSuspense,
@@ -255,9 +256,12 @@ import {
   getSuspendedCache,
   pushTransition,
   getOffscreenDeferredCache,
-  getSuspendedTransitions,
+  getPendingTransitions,
 } from './ReactFiberTransition.new';
-import {pushTracingMarker} from './ReactFiberTracingMarkerComponent.new';
+import {
+  getTracingMarkers,
+  pushTracingMarker,
+} from './ReactFiberTracingMarkerComponent.new';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 
@@ -891,6 +895,20 @@ function updateTracingMarkerComponent(
     return null;
   }
 
+  // TODO: (luna) Only update the tracing marker if it's newly rendered or it's name changed.
+  // A tracing marker is only associated with the transitions that rendered
+  // or updated it, so we can create a new set of transitions each time
+  if (current === null) {
+    const currentTransitions = getPendingTransitions();
+    if (currentTransitions !== null) {
+      const markerInstance: TracingMarkerInstance = {
+        transitions: new Set(currentTransitions),
+        pendingSuspenseBoundaries: new Map(),
+      };
+      workInProgress.stateNode = markerInstance;
+    }
+  }
+
   pushTracingMarker(workInProgress);
   const nextChildren = workInProgress.pendingProps.children;
   reconcileChildren(current, workInProgress, nextChildren, renderLanes);
@@ -2093,10 +2111,13 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
       );
       workInProgress.memoizedState = SUSPENDED_MARKER;
       if (enableTransitionTracing) {
-        const currentTransitions = getSuspendedTransitions();
+        const currentTransitions = getPendingTransitions();
         if (currentTransitions !== null) {
+          // If there are no transitions, we don't need to keep track of tracing markers
+          const currentTracingMarkers = getTracingMarkers();
           const primaryChildUpdateQueue: OffscreenQueue = {
             transitions: currentTransitions,
+            tracingMarkers: currentTracingMarkers,
           };
           primaryChildFragment.updateQueue = primaryChildUpdateQueue;
         }
@@ -2177,10 +2198,12 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
           ? mountSuspenseOffscreenState(renderLanes)
           : updateSuspenseOffscreenState(prevOffscreenState, renderLanes);
       if (enableTransitionTracing) {
-        const currentTransitions = getSuspendedTransitions();
+        const currentTransitions = getPendingTransitions();
         if (currentTransitions !== null) {
+          const currentTracingMarkers = getTracingMarkers();
           const primaryChildUpdateQueue: OffscreenQueue = {
             transitions: currentTransitions,
+            tracingMarkers: currentTracingMarkers,
           };
           primaryChildFragment.updateQueue = primaryChildUpdateQueue;
         }

packages/react-reconciler/src/ReactFiberBeginWork.old.js

@@ -35,6 +35,7 @@ import type {
 } from './ReactFiberCacheComponent.old';
 import type {UpdateQueue} from './ReactFiberClassUpdateQueue.old';
 import type {RootState} from './ReactFiberRoot.old';
+import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.old';
 import {
   enableSuspenseAvoidThisFallback,
   enableCPUSuspense,
@@ -255,9 +256,12 @@ import {
   getSuspendedCache,
   pushTransition,
   getOffscreenDeferredCache,
-  getSuspendedTransitions,
+  getPendingTransitions,
 } from './ReactFiberTransition.old';
-import {pushTracingMarker} from './ReactFiberTracingMarkerComponent.old';
+import {
+  getTracingMarkers,
+  pushTracingMarker,
+} from './ReactFiberTracingMarkerComponent.old';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 
@@ -891,6 +895,20 @@ function updateTracingMarkerComponent(
     return null;
   }
 
+  // TODO: (luna) Only update the tracing marker if it's newly rendered or it's name changed.
+  // A tracing marker is only associated with the transitions that rendered
+  // or updated it, so we can create a new set of transitions each time
+  if (current === null) {
+    const currentTransitions = getPendingTransitions();
+    if (currentTransitions !== null) {
+      const markerInstance: TracingMarkerInstance = {
+        transitions: new Set(currentTransitions),
+        pendingSuspenseBoundaries: new Map(),
+      };
+      workInProgress.stateNode = markerInstance;
+    }
+  }
+
   pushTracingMarker(workInProgress);
   const nextChildren = workInProgress.pendingProps.children;
   reconcileChildren(current, workInProgress, nextChildren, renderLanes);
@@ -2093,10 +2111,13 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
       );
       workInProgress.memoizedState = SUSPENDED_MARKER;
       if (enableTransitionTracing) {
-        const currentTransitions = getSuspendedTransitions();
+        const currentTransitions = getPendingTransitions();
         if (currentTransitions !== null) {
+          // If there are no transitions, we don't need to keep track of tracing markers
+          const currentTracingMarkers = getTracingMarkers();
           const primaryChildUpdateQueue: OffscreenQueue = {
             transitions: currentTransitions,
+            tracingMarkers: currentTracingMarkers,
           };
           primaryChildFragment.updateQueue = primaryChildUpdateQueue;
         }
@@ -2177,10 +2198,12 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
           ? mountSuspenseOffscreenState(renderLanes)
           : updateSuspenseOffscreenState(prevOffscreenState, renderLanes);
       if (enableTransitionTracing) {
-        const currentTransitions = getSuspendedTransitions();
+        const currentTransitions = getPendingTransitions();
         if (currentTransitions !== null) {
+          const currentTracingMarkers = getTracingMarkers();
           const primaryChildUpdateQueue: OffscreenQueue = {
             transitions: currentTransitions,
+            tracingMarkers: currentTracingMarkers,
           };
           primaryChildFragment.updateQueue = primaryChildUpdateQueue;
         }

packages/react-reconciler/src/ReactFiberCommitWork.new.js

@@ -138,6 +138,7 @@ import {
   restorePendingUpdaters,
   addTransitionStartCallbackToPendingTransition,
   addTransitionCompleteCallbackToPendingTransition,
+  addMarkerCompleteCallbackToPendingTransition,
   setIsRunningInsertionEffect,
 } from './ReactFiberWorkLoop.new';
 import {
@@ -2910,6 +2911,7 @@ function commitPassiveMountOnFiber(
               instance.transitions = prevTransitions = new Set();
             }
 
+            // TODO(luna): Combine the root code with the tracing marker code
             if (transitions !== null) {
               transitions.forEach(transition => {
                 // Add all the transitions saved in the update queue during
@@ -2931,6 +2933,23 @@ function commitPassiveMountOnFiber(
                 }
               });
             }
+
+            const tracingMarkers = queue.tracingMarkers;
+            if (tracingMarkers !== null) {
+              tracingMarkers.forEach(marker => {
+                const markerInstance = marker.stateNode;
+                // There should only be a few tracing marker transitions because
+                // they should be only associated with the transition that
+                // caused them
+                markerInstance.transitions.forEach(transition => {
+                  if (instance.transitions.has(transition)) {
+                    instance.pendingMarkers.add(
+                      markerInstance.pendingSuspenseBoundaries,
+                    );
+                  }
+                });
+              });
+            }
           }
 
           commitTransitionProgress(finishedWork);
@@ -2967,6 +2986,28 @@ function commitPassiveMountOnFiber(
       }
       break;
     }
+    case TracingMarkerComponent: {
+      if (enableTransitionTracing) {
+        // Get the transitions that were initiatized during the render
+        // and add a start transition callback for each of them
+        const instance = finishedWork.stateNode;
+        if (
+          instance.pendingSuspenseBoundaries === null ||
+          instance.pendingSuspenseBoundaries.size === 0
+        ) {
+          instance.transitions.forEach(transition => {
+            addMarkerCompleteCallbackToPendingTransition({
+              transitionName: transition.name,
+              startTime: transition.startTime,
+              markerName: finishedWork.memoizedProps.name,
+            });
+          });
+          instance.transitions = null;
+          instance.pendingSuspenseBoundaries = null;
+        }
+      }
+      break;
+    }
   }
 }
 

packages/react-reconciler/src/ReactFiberCommitWork.old.js

@@ -138,6 +138,7 @@ import {
   restorePendingUpdaters,
   addTransitionStartCallbackToPendingTransition,
   addTransitionCompleteCallbackToPendingTransition,
+  addMarkerCompleteCallbackToPendingTransition,
   setIsRunningInsertionEffect,
 } from './ReactFiberWorkLoop.old';
 import {
@@ -2910,6 +2911,7 @@ function commitPassiveMountOnFiber(
               instance.transitions = prevTransitions = new Set();
             }
 
+            // TODO(luna): Combine the root code with the tracing marker code
             if (transitions !== null) {
               transitions.forEach(transition => {
                 // Add all the transitions saved in the update queue during
@@ -2931,6 +2933,23 @@ function commitPassiveMountOnFiber(
                 }
               });
             }
+
+            const tracingMarkers = queue.tracingMarkers;
+            if (tracingMarkers !== null) {
+              tracingMarkers.forEach(marker => {
+                const markerInstance = marker.stateNode;
+                // There should only be a few tracing marker transitions because
+                // they should be only associated with the transition that
+                // caused them
+                markerInstance.transitions.forEach(transition => {
+                  if (instance.transitions.has(transition)) {
+                    instance.pendingMarkers.add(
+                      markerInstance.pendingSuspenseBoundaries,
+                    );
+                  }
+                });
+              });
+            }
           }
 
           commitTransitionProgress(finishedWork);
@@ -2967,6 +2986,28 @@ function commitPassiveMountOnFiber(
       }
       break;
     }
+    case TracingMarkerComponent: {
+      if (enableTransitionTracing) {
+        // Get the transitions that were initiatized during the render
+        // and add a start transition callback for each of them
+        const instance = finishedWork.stateNode;
+        if (
+          instance.pendingSuspenseBoundaries === null ||
+          instance.pendingSuspenseBoundaries.size === 0
+        ) {
+          instance.transitions.forEach(transition => {
+            addMarkerCompleteCallbackToPendingTransition({
+              transitionName: transition.name,
+              startTime: transition.startTime,
+              markerName: finishedWork.memoizedProps.name,
+            });
+          });
+          instance.transitions = null;
+          instance.pendingSuspenseBoundaries = null;
+        }
+      }
+      break;
+    }
   }
 }
 

packages/react-reconciler/src/ReactFiberCompleteWork.new.js

@@ -1581,9 +1581,18 @@ function completeWork(
     }
     case TracingMarkerComponent: {
       if (enableTransitionTracing) {
-        // Bubble subtree flags before so we can set the flag property
         popTracingMarker(workInProgress);
         bubbleProperties(workInProgress);
+
+        if (
+          current === null ||
+          (workInProgress.subtreeFlags & Visibility) !== NoFlags
+        ) {
+          // If any of our suspense children toggle visibility, this means that
+          // the pending boundaries array needs to be updated, which we only
+          // do in the passive phase.
+          workInProgress.flags |= Passive;
+        }
       }
       return null;
     }

packages/react-reconciler/src/ReactFiberCompleteWork.old.js

@@ -1581,9 +1581,18 @@ function completeWork(
     }
     case TracingMarkerComponent: {
       if (enableTransitionTracing) {
-        // Bubble subtree flags before so we can set the flag property
         popTracingMarker(workInProgress);
         bubbleProperties(workInProgress);
+
+        if (
+          current === null ||
+          (workInProgress.subtreeFlags & Visibility) !== NoFlags
+        ) {
+          // If any of our suspense children toggle visibility, this means that
+          // the pending boundaries array needs to be updated, which we only
+          // do in the passive phase.
+          workInProgress.flags |= Passive;
+        }
       }
       return null;
     }

packages/react-reconciler/src/ReactFiberOffscreenComponent.js

@@ -8,6 +8,7 @@
  */
 
 import type {ReactNodeList, OffscreenMode} from 'shared/ReactTypes';
+import type {Fiber} from './ReactInternalTypes';
 import type {Lanes} from './ReactFiberLane.old';
 import type {SpawnedCachePool} from './ReactFiberCacheComponent.new';
 import type {
@@ -38,6 +39,7 @@ export type OffscreenState = {|
 
 export type OffscreenQueue = {|
   transitions: Array<Transition> | null,
+  tracingMarkers: Array<Fiber> | null,
 |} | null;
 
 export type OffscreenInstance = {|

packages/react-reconciler/src/ReactFiberTracingMarkerComponent.new.js

@@ -39,6 +39,11 @@ export type BatchConfigTransition = {
   _updatedFibers?: Set<Fiber>,
 };
 
+export type TracingMarkerInstance = {|
+  pendingSuspenseBoundaries: PendingSuspenseBoundaries | null,
+  transitions: Set<Transition> | null,
+|} | null;
+
 export type PendingSuspenseBoundaries = Map<OffscreenInstance, SuspenseInfo>;
 
 export function processTransitionCallbacks(