[Transition Tracing] Add transition to OffscreenState and pendingSuspenseBoundaries to RootState (#24340)

In this PR we:

Add transitions boilerplate to the OffscreenState. The transitions field will be null on initiation. During the commit phase, if there are any new transitions, we will add any new transitions (either as a result of a transition occurring or a parent suspense boundary completing) to the transitions field. Once the suspense boundary resolves, we no longer need to store the transitions on the boundary, so we can put this field on the Offscreen memoized state
Add pendingSuspenseBoundaries boilerplate to the RootState. This field starts as null. During the commit phase, if a suspense boundary has either gone from fallback to resolved or from resolved to fallback, we will create a new Map if there isn't one, and if there is, we will add (if the boundary is a fallback) or remove the suspense boundary (if the boundary has resolved) from the map.
Add an optional name field to the Suspense boundary

Author: lunaruan

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

@@ -236,7 +236,6 @@ import {
   markSkippedUpdateLanes,
   getWorkInProgressRoot,
   pushRenderLanes,
-  getWorkInProgressTransitions,
 } from './ReactFiberWorkLoop.new';
 import {setWorkInProgressVersion} from './ReactMutableSource.new';
 import {pushCacheProvider, CacheContext} from './ReactFiberCacheComponent.new';
@@ -652,9 +651,11 @@ function updateOffscreenComponent(
     // Rendering a hidden tree.
     if ((workInProgress.mode & ConcurrentMode) === NoMode) {
       // In legacy sync mode, don't defer the subtree. Render it now.
+      // TODO: Consider how Offscreen should work with transitions in the future
       const nextState: OffscreenState = {
         baseLanes: NoLanes,
         cachePool: null,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       if (enableCache) {
@@ -688,6 +689,7 @@ function updateOffscreenComponent(
       const nextState: OffscreenState = {
         baseLanes: nextBaseLanes,
         cachePool: spawnedCachePool,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       workInProgress.updateQueue = null;
@@ -723,6 +725,7 @@ function updateOffscreenComponent(
       const nextState: OffscreenState = {
         baseLanes: NoLanes,
         cachePool: null,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       // Push the lanes that were skipped when we bailed out.
@@ -1345,13 +1348,6 @@ function updateHostRoot(current, workInProgress, renderLanes) {
     }
   }
 
-  if (enableTransitionTracing) {
-    // FIXME: Slipped past code review. This is not a safe mutation:
-    // workInProgress.memoizedState is a shared object. Need to fix before
-    // rolling out the Transition Tracing experiment.
-    workInProgress.memoizedState.transitions = getWorkInProgressTransitions();
-  }
-
   // Caution: React DevTools currently depends on this property
   // being called "element".
   const nextChildren = nextState.element;
@@ -1365,6 +1361,7 @@ function updateHostRoot(current, workInProgress, renderLanes) {
       element: nextChildren,
       isDehydrated: false,
       cache: nextState.cache,
+      pendingSuspenseBoundaries: nextState.pendingSuspenseBoundaries,
       transitions: nextState.transitions,
     };
     const updateQueue: UpdateQueue<RootState> = (workInProgress.updateQueue: any);
@@ -1982,6 +1979,7 @@ function mountSuspenseOffscreenState(renderLanes: Lanes): OffscreenState {
   return {
     baseLanes: renderLanes,
     cachePool: getSuspendedCache(),
+    transitions: null,
   };
 }
 
@@ -2016,6 +2014,7 @@ function updateSuspenseOffscreenState(
   return {
     baseLanes: mergeLanes(prevOffscreenState.baseLanes, renderLanes),
     cachePool,
+    transitions: prevOffscreenState.transitions,
   };
 }
 
@@ -3582,9 +3581,6 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
         const cache: Cache = current.memoizedState.cache;
         pushCacheProvider(workInProgress, cache);
       }
-      if (enableTransitionTracing) {
-        workInProgress.memoizedState.transitions = getWorkInProgressTransitions();
-      }
       resetHydrationState();
       break;
     case HostComponent:

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

@@ -236,7 +236,6 @@ import {
   markSkippedUpdateLanes,
   getWorkInProgressRoot,
   pushRenderLanes,
-  getWorkInProgressTransitions,
 } from './ReactFiberWorkLoop.old';
 import {setWorkInProgressVersion} from './ReactMutableSource.old';
 import {pushCacheProvider, CacheContext} from './ReactFiberCacheComponent.old';
@@ -652,9 +651,11 @@ function updateOffscreenComponent(
     // Rendering a hidden tree.
     if ((workInProgress.mode & ConcurrentMode) === NoMode) {
       // In legacy sync mode, don't defer the subtree. Render it now.
+      // TODO: Consider how Offscreen should work with transitions in the future
       const nextState: OffscreenState = {
         baseLanes: NoLanes,
         cachePool: null,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       if (enableCache) {
@@ -688,6 +689,7 @@ function updateOffscreenComponent(
       const nextState: OffscreenState = {
         baseLanes: nextBaseLanes,
         cachePool: spawnedCachePool,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       workInProgress.updateQueue = null;
@@ -723,6 +725,7 @@ function updateOffscreenComponent(
       const nextState: OffscreenState = {
         baseLanes: NoLanes,
         cachePool: null,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       // Push the lanes that were skipped when we bailed out.
@@ -1345,13 +1348,6 @@ function updateHostRoot(current, workInProgress, renderLanes) {
     }
   }
 
-  if (enableTransitionTracing) {
-    // FIXME: Slipped past code review. This is not a safe mutation:
-    // workInProgress.memoizedState is a shared object. Need to fix before
-    // rolling out the Transition Tracing experiment.
-    workInProgress.memoizedState.transitions = getWorkInProgressTransitions();
-  }
-
   // Caution: React DevTools currently depends on this property
   // being called "element".
   const nextChildren = nextState.element;
@@ -1365,6 +1361,7 @@ function updateHostRoot(current, workInProgress, renderLanes) {
       element: nextChildren,
       isDehydrated: false,
       cache: nextState.cache,
+      pendingSuspenseBoundaries: nextState.pendingSuspenseBoundaries,
       transitions: nextState.transitions,
     };
     const updateQueue: UpdateQueue<RootState> = (workInProgress.updateQueue: any);
@@ -1982,6 +1979,7 @@ function mountSuspenseOffscreenState(renderLanes: Lanes): OffscreenState {
   return {
     baseLanes: renderLanes,
     cachePool: getSuspendedCache(),
+    transitions: null,
   };
 }
 
@@ -2016,6 +2014,7 @@ function updateSuspenseOffscreenState(
   return {
     baseLanes: mergeLanes(prevOffscreenState.baseLanes, renderLanes),
     cachePool,
+    transitions: prevOffscreenState.transitions,
   };
 }
 
@@ -3582,9 +3581,6 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
         const cache: Cache = current.memoizedState.cache;
         pushCacheProvider(workInProgress, cache);
       }
-      if (enableTransitionTracing) {
-        workInProgress.memoizedState.transitions = getWorkInProgressTransitions();
-      }
       resetHydrationState();
       break;
     case HostComponent:

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

@@ -2725,9 +2725,8 @@ function commitPassiveMountOnFiber(
       }
 
       if (enableTransitionTracing) {
-        const transitions = finishedWork.memoizedState.transitions;
-        if (transitions !== null) {
-          transitions.forEach(transition => {
+        if (committedTransitions !== null) {
+          committedTransitions.forEach(transition => {
             // TODO(luna) Do we want to log TransitionStart in the startTransition callback instead?
             addTransitionStartCallbackToPendingTransition({
               transitionName: transition.name,

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

@@ -2725,9 +2725,8 @@ function commitPassiveMountOnFiber(
       }
 
       if (enableTransitionTracing) {
-        const transitions = finishedWork.memoizedState.transitions;
-        if (transitions !== null) {
-          transitions.forEach(transition => {
+        if (committedTransitions !== null) {
+          committedTransitions.forEach(transition => {
             // TODO(luna) Do we want to log TransitionStart in the startTransition callback instead?
             addTransitionStartCallbackToPendingTransition({
               transitionName: transition.name,

packages/react-reconciler/src/ReactFiberOffscreenComponent.js

@@ -10,7 +10,7 @@
 import type {ReactNodeList, OffscreenMode} from 'shared/ReactTypes';
 import type {Lanes} from './ReactFiberLane.old';
 import type {SpawnedCachePool} from './ReactFiberCacheComponent.new';
-
+import type {Transition} from './ReactFiberTracingMarkerComponent.new';
 export type OffscreenProps = {|
   // TODO: Pick an API before exposing the Offscreen type. I've chosen an enum
   // for now, since we might have multiple variants. For example, hiding the
@@ -30,6 +30,7 @@ export type OffscreenState = {|
   // order to unhide the component.
   baseLanes: Lanes,
   cachePool: SpawnedCachePool | null,
+  transitions: Set<Transition> | null,
 |};
 
 export type OffscreenInstance = {};

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

@@ -15,7 +15,10 @@ import type {
 } from './ReactInternalTypes';
 import type {RootTag} from './ReactRootTags';
 import type {Cache} from './ReactFiberCacheComponent.new';
-import type {Transition} from './ReactFiberTracingMarkerComponent.new';
+import type {
+  PendingSuspenseBoundaries,
+  Transition,
+} from './ReactFiberTracingMarkerComponent.new';
 
 import {noTimeout, supportsHydration} from './ReactFiberHostConfig';
 import {createHostRootFiber} from './ReactFiber.new';
@@ -42,7 +45,8 @@ export type RootState = {
   element: any,
   isDehydrated: boolean,
   cache: Cache,
-  transitions: Array<Transition> | null,
+  pendingSuspenseBoundaries: PendingSuspenseBoundaries | null,
+  transitions: Set<Transition> | null,
 };
 
 function FiberRootNode(
@@ -184,6 +188,7 @@ export function createFiberRoot(
       isDehydrated: hydrate,
       cache: initialCache,
       transitions: null,
+      pendingSuspenseBoundaries: null,
     };
     uninitializedFiber.memoizedState = initialState;
   } else {
@@ -192,6 +197,7 @@ export function createFiberRoot(
       isDehydrated: hydrate,
       cache: (null: any), // not enabled yet
       transitions: null,
+      pendingSuspenseBoundaries: null,
     };
     uninitializedFiber.memoizedState = initialState;
   }

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

@@ -15,7 +15,10 @@ import type {
 } from './ReactInternalTypes';
 import type {RootTag} from './ReactRootTags';
 import type {Cache} from './ReactFiberCacheComponent.old';
-import type {Transition} from './ReactFiberTracingMarkerComponent.old';
+import type {
+  PendingSuspenseBoundaries,
+  Transition,
+} from './ReactFiberTracingMarkerComponent.old';
 
 import {noTimeout, supportsHydration} from './ReactFiberHostConfig';
 import {createHostRootFiber} from './ReactFiber.old';
@@ -42,7 +45,8 @@ export type RootState = {
   element: any,
   isDehydrated: boolean,
   cache: Cache,
-  transitions: Array<Transition> | null,
+  pendingSuspenseBoundaries: PendingSuspenseBoundaries | null,
+  transitions: Set<Transition> | null,
 };
 
 function FiberRootNode(
@@ -184,6 +188,7 @@ export function createFiberRoot(
       isDehydrated: hydrate,
       cache: initialCache,
       transitions: null,
+      pendingSuspenseBoundaries: null,
     };
     uninitializedFiber.memoizedState = initialState;
   } else {
@@ -192,6 +197,7 @@ export function createFiberRoot(
       isDehydrated: hydrate,
       cache: (null: any), // not enabled yet
       transitions: null,
+      pendingSuspenseBoundaries: null,
     };
     uninitializedFiber.memoizedState = initialState;
   }

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

@@ -29,6 +29,7 @@ export type SuspenseProps = {|
   suspenseCallback?: (Set<Wakeable> | null) => mixed,
 
   unstable_expectedLoadTime?: number,
+  unstable_name?: string,
 |};
 
 // A null SuspenseState represents an unsuspended normal Suspense boundary.

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

@@ -29,6 +29,7 @@ export type SuspenseProps = {|
   suspenseCallback?: (Set<Wakeable> | null) => mixed,
 
   unstable_expectedLoadTime?: number,
+  unstable_name?: string,
 |};
 
 // A null SuspenseState represents an unsuspended normal Suspense boundary.

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

@@ -7,8 +7,9 @@
  * @flow
  */
 
-import type {TransitionTracingCallbacks} from './ReactInternalTypes';
-import type {Fiber} from 'react-reconciler/src/ReactInternalTypes';
+import type {TransitionTracingCallbacks, Fiber} from './ReactInternalTypes';
+import type {OffscreenInstance} from './ReactFiberOffscreenComponent';
+
 import {enableTransitionTracing} from 'shared/ReactFeatureFlags';
 
 export type SuspenseInfo = {name: string | null};
@@ -34,10 +35,7 @@ export type BatchConfigTransition = {
   _updatedFibers?: Set<Fiber>,
 };
 
-export type TransitionCallback = 0 | 1;
-
-export const TransitionStart = 0;
-export const TransitionComplete = 1;
+export type PendingSuspenseBoundaries = Map<OffscreenInstance, SuspenseInfo>;
 
 export function processTransitionCallbacks(
   pendingTransitions: PendingTransitionCallbacks,

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

@@ -7,8 +7,9 @@
  * @flow
  */
 
-import type {TransitionTracingCallbacks} from './ReactInternalTypes';
-import type {Fiber} from 'react-reconciler/src/ReactInternalTypes';
+import type {TransitionTracingCallbacks, Fiber} from './ReactInternalTypes';
+import type {OffscreenInstance} from './ReactFiberOffscreenComponent';
+
 import {enableTransitionTracing} from 'shared/ReactFeatureFlags';
 
 export type SuspenseInfo = {name: string | null};
@@ -34,10 +35,7 @@ export type BatchConfigTransition = {
   _updatedFibers?: Set<Fiber>,
 };
 
-export type TransitionCallback = 0 | 1;
-
-export const TransitionStart = 0;
-export const TransitionComplete = 1;
+export type PendingSuspenseBoundaries = Map<OffscreenInstance, SuspenseInfo>;
 
 export function processTransitionCallbacks(
   pendingTransitions: PendingTransitionCallbacks,