[Transition Tracing] Add Tracing Marker Stack

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

@@ -258,6 +258,7 @@ import {
   getOffscreenDeferredCache,
   getSuspendedTransitions,
 } from './ReactFiberTransition.new';
+import {pushTracingMarker} from './ReactFiberTracingMarkerComponent.new';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 
@@ -887,6 +888,7 @@ function updateTracingMarkerComponent(
     return null;
   }
 
+  pushTracingMarker(workInProgress);
   const nextChildren = workInProgress.pendingProps.children;
   reconcileChildren(current, workInProgress, nextChildren, renderLanes);
   return workInProgress.child;
@@ -3678,6 +3680,11 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
       }
       break;
     }
+    case TracingMarkerComponent: {
+      if (enableTransitionTracing) {
+        pushTracingMarker(workInProgress);
+      }
+    }
   }
   return bailoutOnAlreadyFinishedWork(current, workInProgress, renderLanes);
 }

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

@@ -258,6 +258,7 @@ import {
   getOffscreenDeferredCache,
   getSuspendedTransitions,
 } from './ReactFiberTransition.old';
+import {pushTracingMarker} from './ReactFiberTracingMarkerComponent.old';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 
@@ -887,6 +888,7 @@ function updateTracingMarkerComponent(
     return null;
   }
 
+  pushTracingMarker(workInProgress);
   const nextChildren = workInProgress.pendingProps.children;
   reconcileChildren(current, workInProgress, nextChildren, renderLanes);
   return workInProgress.child;
@@ -3678,6 +3680,11 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
       }
       break;
     }
+    case TracingMarkerComponent: {
+      if (enableTransitionTracing) {
+        pushTracingMarker(workInProgress);
+      }
+    }
   }
   return bailoutOnAlreadyFinishedWork(current, workInProgress, renderLanes);
 }

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

@@ -165,6 +165,7 @@ import {transferActualDuration} from './ReactProfilerTimer.new';
 import {popCacheProvider} from './ReactFiberCacheComponent.new';
 import {popTreeContext} from './ReactFiberTreeContext.new';
 import {popRootTransition, popTransition} from './ReactFiberTransition.new';
+import {popTracingMarker} from './ReactFiberTracingMarkerComponent.new';
 
 function markUpdate(workInProgress: Fiber) {
   // Tag the fiber with an update effect. This turns a Placement into
@@ -1585,6 +1586,7 @@ function completeWork(
     case TracingMarkerComponent: {
       if (enableTransitionTracing) {
         // Bubble subtree flags before so we can set the flag property
+        popTracingMarker(workInProgress);
         bubbleProperties(workInProgress);
       }
       return null;

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

@@ -165,6 +165,7 @@ import {transferActualDuration} from './ReactProfilerTimer.old';
 import {popCacheProvider} from './ReactFiberCacheComponent.old';
 import {popTreeContext} from './ReactFiberTreeContext.old';
 import {popRootTransition, popTransition} from './ReactFiberTransition.old';
+import {popTracingMarker} from './ReactFiberTracingMarkerComponent.old';
 
 function markUpdate(workInProgress: Fiber) {
   // Tag the fiber with an update effect. This turns a Placement into
@@ -1585,6 +1586,7 @@ function completeWork(
     case TracingMarkerComponent: {
       if (enableTransitionTracing) {
         // Bubble subtree flags before so we can set the flag property
+        popTracingMarker(workInProgress);
         bubbleProperties(workInProgress);
       }
       return null;

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

@@ -9,8 +9,10 @@
 
 import type {TransitionTracingCallbacks, Fiber} from './ReactInternalTypes';
 import type {OffscreenInstance} from './ReactFiberOffscreenComponent';
+import type {StackCursor} from './ReactFiberStack.new';
 
 import {enableTransitionTracing} from 'shared/ReactFeatureFlags';
+import {createCursor, push, pop} from './ReactFiberStack.new';
 
 export type SuspenseInfo = {name: string | null};
 
@@ -71,3 +73,37 @@ export function processTransitionCallbacks(
     }
   }
 }
+
+// For every tracing marker, store a pointer to it. We will later access it
+// to get the set of suspense boundaries that need to resolve before the
+// tracing marker can be logged as complete
+// This code lives separate from the ReactFiberTransition code because
+// we push and pop on the tracing marker, not the suspense boundary
+const tracingMarkerStack: StackCursor<Array<Fiber> | null> = createCursor(null);
+
+export function pushTracingMarker(workInProgress: Fiber): void {
+  if (enableTransitionTracing) {
+    if (tracingMarkerStack.current === null) {
+      push(tracingMarkerStack, [workInProgress], workInProgress);
+    } else {
+      push(
+        tracingMarkerStack,
+        tracingMarkerStack.current.concat(workInProgress),
+        workInProgress,
+      );
+    }
+  }
+}
+
+export function popTracingMarker(workInProgress: Fiber): void {
+  if (enableTransitionTracing) {
+    pop(tracingMarkerStack, workInProgress);
+  }
+}
+
+export function getTracingMarkers(): Array<Fiber> | null {
+  if (enableTransitionTracing) {
+    return tracingMarkerStack.current;
+  }
+  return null;
+}

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

@@ -9,8 +9,10 @@
 
 import type {TransitionTracingCallbacks, Fiber} from './ReactInternalTypes';
 import type {OffscreenInstance} from './ReactFiberOffscreenComponent';
+import type {StackCursor} from './ReactFiberStack.old';
 
 import {enableTransitionTracing} from 'shared/ReactFeatureFlags';
+import {createCursor, push, pop} from './ReactFiberStack.old';
 
 export type SuspenseInfo = {name: string | null};
 
@@ -71,3 +73,37 @@ export function processTransitionCallbacks(
     }
   }
 }
+
+// For every tracing marker, store a pointer to it. We will later access it
+// to get the set of suspense boundaries that need to resolve before the
+// tracing marker can be logged as complete
+// This code lives separate from the ReactFiberTransition code because
+// we push and pop on the tracing marker, not the suspense boundary
+const tracingMarkerStack: StackCursor<Array<Fiber> | null> = createCursor(null);
+
+export function pushTracingMarker(workInProgress: Fiber): void {
+  if (enableTransitionTracing) {
+    if (tracingMarkerStack.current === null) {
+      push(tracingMarkerStack, [workInProgress], workInProgress);
+    } else {
+      push(
+        tracingMarkerStack,
+        tracingMarkerStack.current.concat(workInProgress),
+        workInProgress,
+      );
+    }
+  }
+}
+
+export function popTracingMarker(workInProgress: Fiber): void {
+  if (enableTransitionTracing) {
+    pop(tracingMarkerStack, workInProgress);
+  }
+}
+
+export function getTracingMarkers(): Array<Fiber> | null {
+  if (enableTransitionTracing) {
+    return tracingMarkerStack.current;
+  }
+  return null;
+}

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

@@ -25,10 +25,15 @@ import {
   OffscreenComponent,
   LegacyHiddenComponent,
   CacheComponent,
+  TracingMarkerComponent,
 } from './ReactWorkTags';
 import {DidCapture, NoFlags, ShouldCapture} from './ReactFiberFlags';
 import {NoMode, ProfileMode} from './ReactTypeOfMode';
-import {enableProfilerTimer, enableCache} from 'shared/ReactFeatureFlags';
+import {
+  enableProfilerTimer,
+  enableCache,
+  enableTransitionTracing,
+} from 'shared/ReactFeatureFlags';
 
 import {popHostContainer, popHostContext} from './ReactFiberHostContext.new';
 import {popSuspenseContext} from './ReactFiberSuspenseContext.new';
@@ -44,6 +49,7 @@ import {popCacheProvider} from './ReactFiberCacheComponent.new';
 import {transferActualDuration} from './ReactProfilerTimer.new';
 import {popTreeContext} from './ReactFiberTreeContext.new';
 import {popRootTransition, popTransition} from './ReactFiberTransition.new';
+import {popTracingMarker} from './ReactFiberTracingMarkerComponent.new';
 
 function unwindWork(
   current: Fiber | null,
@@ -154,6 +160,11 @@ function unwindWork(
         popCacheProvider(workInProgress, cache);
       }
       return null;
+    case TracingMarkerComponent:
+      if (enableTransitionTracing) {
+        popTracingMarker(workInProgress);
+      }
+      return null;
     default:
       return null;
   }
@@ -217,6 +228,11 @@ function unwindInterruptedWork(
         popCacheProvider(interruptedWork, cache);
       }
       break;
+    case TracingMarkerComponent:
+      if (enableTransitionTracing) {
+        popTracingMarker(interruptedWork);
+      }
+      break;
     default:
       break;
   }

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

@@ -25,10 +25,15 @@ import {
   OffscreenComponent,
   LegacyHiddenComponent,
   CacheComponent,
+  TracingMarkerComponent,
 } from './ReactWorkTags';
 import {DidCapture, NoFlags, ShouldCapture} from './ReactFiberFlags';
 import {NoMode, ProfileMode} from './ReactTypeOfMode';
-import {enableProfilerTimer, enableCache} from 'shared/ReactFeatureFlags';
+import {
+  enableProfilerTimer,
+  enableCache,
+  enableTransitionTracing,
+} from 'shared/ReactFeatureFlags';
 
 import {popHostContainer, popHostContext} from './ReactFiberHostContext.old';
 import {popSuspenseContext} from './ReactFiberSuspenseContext.old';
@@ -44,6 +49,7 @@ import {popCacheProvider} from './ReactFiberCacheComponent.old';
 import {transferActualDuration} from './ReactProfilerTimer.old';
 import {popTreeContext} from './ReactFiberTreeContext.old';
 import {popRootTransition, popTransition} from './ReactFiberTransition.old';
+import {popTracingMarker} from './ReactFiberTracingMarkerComponent.old';
 
 function unwindWork(
   current: Fiber | null,
@@ -154,6 +160,11 @@ function unwindWork(
         popCacheProvider(workInProgress, cache);
       }
       return null;
+    case TracingMarkerComponent:
+      if (enableTransitionTracing) {
+        popTracingMarker(workInProgress);
+      }
+      return null;
     default:
       return null;
   }
@@ -217,6 +228,11 @@ function unwindInterruptedWork(
         popCacheProvider(interruptedWork, cache);
       }
       break;
+    case TracingMarkerComponent:
+      if (enableTransitionTracing) {
+        popTracingMarker(interruptedWork);
+      }
+      break;
     default:
       break;
   }