Repository: facebook/react
Pull request: [DevTools][Transition Tracing] Added support for Suspense Boundaries

This PR:
* During the passive effect complete phase for Offscreen, we add all the transitions that were added to the update queue in the render phase to the transitions set on the memoziedState. We also add the stateNode for the Offscreen component to the root pendingSuspenseBoundaries map if the suspense boundary has gone from shown to fallback. We remove the stateNode if the boundary goes from fallback to shown.
* During the passive effect complete phase for the HostRoot, for each transition that was initiated during this commit, we add a pending transitionStart callback. We also add them to the transition set on the memoizedState for the HostRoot. If the root pendingSuspenseBoundaries is empty, we add a pending transitionComplete callback.

Changed files:
filename: packages/react-reconciler/src/ReactFiberCommitWork.new.js
patch: @@ -22,7 +22,10 @@ import type {SuspenseState} from './ReactFiberSuspenseComponent.new';
 import type {UpdateQueue} from './ReactUpdateQueue.new';
 import type {FunctionComponentUpdateQueue} from './ReactFiberHooks.new';
 import type {Wakeable} from 'shared/ReactTypes';
-import type {OffscreenState} from './ReactFiberOffscreenComponent';
+import type {
+  OffscreenState,
+  OffscreenInstance,
+} from './ReactFiberOffscreenComponent';
 import type {HookFlags} from './ReactHookEffectTags';
 import type {Cache} from './ReactFiberCacheComponent.new';
 import type {RootState} from './ReactFiberRoot.new';
@@ -62,6 +65,7 @@ import {
   OffscreenComponent,
   LegacyHiddenComponent,
   CacheComponent,
+  TracingMarkerComponent,
 } from './ReactWorkTags';
 import {detachDeletedInstance} from './ReactFiberHostConfig';
 import {
@@ -1001,7 +1005,8 @@ function commitLayoutEffectOnFiber(
       case IncompleteClassComponent:
       case ScopeComponent:
       case OffscreenComponent:
-      case LegacyHiddenComponent: {
+      case LegacyHiddenComponent:
+      case TracingMarkerComponent: {
         break;
       }
 
@@ -1066,6 +1071,89 @@ function reappearLayoutEffectsOnFiber(node: Fiber) {
   }
 }
 
+function commitTransitionProgress(
+  finishedRoot: FiberRoot,
+  offscreenFiber: Fiber,
+) {
+  if (enableTransitionTracing) {
+    // This function adds suspense boundaries to the root
+    // or tracing marker's pendingSuspenseBoundaries map.
+    // When a suspense boundary goes from a resolved to a fallback
+    // state we add the boundary to the map, and when it goes from
+    // a fallback to a resolved state, we remove the boundary from
+    // the map.
+
+    // We use stateNode on the Offscreen component as a stable object
+    // that doesnt change from render to render. This way we can
+    // distinguish between different Offscreen instances (vs. the same
+    // Offscreen instance with different fibers)
+    const offscreenInstance: OffscreenInstance = offscreenFiber.stateNode;
+
+    let prevState: SuspenseState | null = null;
+    const previousFiber = offscreenFiber.alternate;
+    if (previousFiber !== null && previousFiber.memoizedState !== null) {
+      prevState = previousFiber.memoizedState;
+    }
+    const nextState: SuspenseState | null = offscreenFiber.memoizedState;
+
+    const wasHidden = prevState !== null;
+    const isHidden = nextState !== null;
+
+    const rootState: RootState = finishedRoot.current.memoizedState;
+    // TODO(luna) move pendingSuspenseBoundaries and transitions from
+    // HostRoot fiber to FiberRoot
+    const rootPendingBoundaries = rootState.pendingSuspenseBoundaries;
+    const rootTransitions = rootState.transitions;
+
+    // If there is a name on the suspense boundary, store that in
+    // the pending boundaries.
+    let name = null;
+    const parent = offscreenFiber.return;
+    if (
+      parent !== null &&
+      parent.tag === SuspenseComponent &&
+      parent.memoizedProps.unstable_name
+    ) {
+      name = parent.memoizedProps.unstable_name;
+    }
+
+    if (rootPendingBoundaries !== null) {
+      if (previousFiber === null) {
+        // Initial mount
+        if (isHidden) {
+          rootPendingBoundaries.set(offscreenInstance, {
+            name,
+          });
+        }
+      } else {
+        if (wasHidden && !isHidden) {
+          // The suspense boundary went from hidden to visible. Remove
+          // the boundary from the pending suspense boundaries set
+          // if it's there
+          if (rootPendingBoundaries.has(offscreenInstance)) {
+            rootPendingBoundaries.delete(offscreenInstance);
+
+            if (rootPendingBoundaries.size === 0 && rootTransitions !== null) {
+              rootTransitions.forEach(transition => {
+                addTransitionCompleteCallbackToPendingTransition({
+                  transitionName: transition.name,
+                  startTime: transition.startTime,
+                });
+              });
+            }
+          }
+        } else if (!wasHidden && isHidden) {
+          // The suspense boundaries was just hidden. Add the boundary
+          // to the pending boundary set if it's there
+          rootPendingBoundaries.set(offscreenInstance, {
+            name,
+          });
+        }
+      }
+    }
+  }
+}
+
 function hideOrUnhideAllChildren(finishedWork, isHidden) {
   // Only hide or unhide the top-most host nodes.
   let hostSubtreeRoot = null;
@@ -2747,22 +2835,48 @@ function commitPassiveMountOnFiber(
       }
 
       if (enableTransitionTracing) {
+        // Get the transitions that were initiatized during the render
+        // and add a start transition callback for each of them
+        const state = finishedWork.memoizedState;
+        // TODO Since it's a mutable field, this should live on the FiberRoot
+        if (state.transitions === null) {
+          state.transitions = new Set([]);
+        }
+        const pendingTransitions = state.transitions;
+        const pendingSuspenseBoundaries = state.pendingSuspenseBoundaries;
+
+        // Initial render
         if (committedTransitions !== null) {
           committedTransitions.forEach(transition => {
-            // TODO(luna) Do we want to log TransitionStart in the startTransition callback instead?
             addTransitionStartCallbackToPendingTransition({
               transitionName: transition.name,
               startTime: transition.startTime,
             });
+            pendingTransitions.add(transition);
+          });
 
-            addTransitionCompleteCallbackToPendingTransition({
-              transitionName: transition.name,
-              startTime: transition.startTime,
+          if (
+            pendingSuspenseBoundaries === null ||
+            pendingSuspenseBoundaries.size === 0
+          ) {
+            pendingTransitions.forEach(transition => {
+              addTransitionCompleteCallbackToPendingTransition({
+                transitionName: transition.name,
+                startTime: transition.startTime,
+              });
             });
-          });
+          }
 
           clearTransitionsForLanes(finishedRoot, committedLanes);
-          finishedWork.memoizedState.transitions = null;
+        }
+
+        // If there are no more pending suspense boundaries we
+        // clear the transitions because they are all complete.
+        if (
+          pendingSuspenseBoundaries === null ||
+          pendingSuspenseBoundaries.size === 0
+        ) {
+          state.transitions = null;
         }
       }
       break;
@@ -2800,9 +2914,44 @@ function commitPassiveMountOnFiber(
       }
 
       if (enableTransitionTracing) {
-        // TODO: Add code to actually process the update queue
+        const isFallback = finishedWork.memoizedState;
+        const queue = (finishedWork.updateQueue: any);
+        const rootMemoizedState = finishedRoot.current.memoizedState;
+
+        if (queue !== null) {
+          // We have one instance of the pendingSuspenseBoundaries map.
+          // We only need one because we update it during the commit phase.
+          // We instantiate a new Map if we haven't already
+          if (rootMemoizedState.pendingSuspenseBoundaries === null) {
+            rootMemoizedState.pendingSuspenseBoundaries = new Map();
+          }
+
+          if (isFallback) {
+            const transitions = queue.transitions;
+            let prevTransitions = finishedWork.memoizedState.transitions;
+            // Add all the transitions saved in the update queue during
+            // the render phase (ie the transitions associated with this boundary)
+            // into the transitions set.
+            if (transitions !== null) {
+              if (prevTransitions === null) {
+                // We only have one instance of the transitions set
+                // because we update it only during the commit phase. We
+                // will create the set on a as needed basis in the commit phase
+                finishedWork.memoizedState.transitions = prevTransitions = new Set();
+              }
+
+              transitions.forEach(transition => {
+                prevTransitions.add(transition);
+              });
+            }
+          }
+        }
+
+        commitTransitionProgress(finishedRoot, finishedWork);
+
         finishedWork.updateQueue = null;
       }
+
       break;
     }
     case CacheComponent: {

filename: packages/react-reconciler/src/ReactFiberCommitWork.old.js
patch: @@ -22,7 +22,10 @@ import type {SuspenseState} from './ReactFiberSuspenseComponent.old';
 import type {UpdateQueue} from './ReactUpdateQueue.old';
 import type {FunctionComponentUpdateQueue} from './ReactFiberHooks.old';
 import type {Wakeable} from 'shared/ReactTypes';
-import type {OffscreenState} from './ReactFiberOffscreenComponent';
+import type {
+  OffscreenState,
+  OffscreenInstance,
+} from './ReactFiberOffscreenComponent';
 import type {HookFlags} from './ReactHookEffectTags';
 import type {Cache} from './ReactFiberCacheComponent.old';
 import type {RootState} from './ReactFiberRoot.old';
@@ -62,6 +65,7 @@ import {
   OffscreenComponent,
   LegacyHiddenComponent,
   CacheComponent,
+  TracingMarkerComponent,
 } from './ReactWorkTags';
 import {detachDeletedInstance} from './ReactFiberHostConfig';
 import {
@@ -1001,7 +1005,8 @@ function commitLayoutEffectOnFiber(
       case IncompleteClassComponent:
       case ScopeComponent:
       case OffscreenComponent:
-      case LegacyHiddenComponent: {
+      case LegacyHiddenComponent:
+      case TracingMarkerComponent: {
         break;
       }
 
@@ -1066,6 +1071,89 @@ function reappearLayoutEffectsOnFiber(node: Fiber) {
   }
 }
 
+function commitTransitionProgress(
+  finishedRoot: FiberRoot,
+  offscreenFiber: Fiber,
+) {
+  if (enableTransitionTracing) {
+    // This function adds suspense boundaries to the root
+    // or tracing marker's pendingSuspenseBoundaries map.
+    // When a suspense boundary goes from a resolved to a fallback
+    // state we add the boundary to the map, and when it goes from
+    // a fallback to a resolved state, we remove the boundary from
+    // the map.
+
+    // We use stateNode on the Offscreen component as a stable object
+    // that doesnt change from render to render. This way we can
+    // distinguish between different Offscreen instances (vs. the same
+    // Offscreen instance with different fibers)
+    const offscreenInstance: OffscreenInstance = offscreenFiber.stateNode;
+
+    let prevState: SuspenseState | null = null;
+    const previousFiber = offscreenFiber.alternate;
+    if (previousFiber !== null && previousFiber.memoizedState !== null) {
+      prevState = previousFiber.memoizedState;
+    }
+    const nextState: SuspenseState | null = offscreenFiber.memoizedState;
+
+    const wasHidden = prevState !== null;
+    const isHidden = nextState !== null;
+
+    const rootState: RootState = finishedRoot.current.memoizedState;
+    // TODO(luna) move pendingSuspenseBoundaries and transitions from
+    // HostRoot fiber to FiberRoot
+    const rootPendingBoundaries = rootState.pendingSuspenseBoundaries;
+    const rootTransitions = rootState.transitions;
+
+    // If there is a name on the suspense boundary, store that in
+    // the pending boundaries.
+    let name = null;
+    const parent = offscreenFiber.return;
+    if (
+      parent !== null &&
+      parent.tag === SuspenseComponent &&
+      parent.memoizedProps.unstable_name
+    ) {
+      name = parent.memoizedProps.unstable_name;
+    }
+
+    if (rootPendingBoundaries !== null) {
+      if (previousFiber === null) {
+        // Initial mount
+        if (isHidden) {
+          rootPendingBoundaries.set(offscreenInstance, {
+            name,
+          });
+        }
+      } else {
+        if (wasHidden && !isHidden) {
+          // The suspense boundary went from hidden to visible. Remove
+          // the boundary from the pending suspense boundaries set
+          // if it's there
+          if (rootPendingBoundaries.has(offscreenInstance)) {
+            rootPendingBoundaries.delete(offscreenInstance);
+
+            if (rootPendingBoundaries.size === 0 && rootTransitions !== null) {
+              rootTransitions.forEach(transition => {
+                addTransitionCompleteCallbackToPendingTransition({
+                  transitionName: transition.name,
+                  startTime: transition.startTime,
+                });
+              });
+            }
+          }
+        } else if (!wasHidden && isHidden) {
+          // The suspense boundaries was just hidden. Add the boundary
+          // to the pending boundary set if it's there
+          rootPendingBoundaries.set(offscreenInstance, {
+            name,
+          });
+        }
+      }
+    }
+  }
+}
+
 function hideOrUnhideAllChildren(finishedWork, isHidden) {
   // Only hide or unhide the top-most host nodes.
   let hostSubtreeRoot = null;
@@ -2747,22 +2835,48 @@ function commitPassiveMountOnFiber(
       }
 
       if (enableTransitionTracing) {
+        // Get the transitions that were initiatized during the render
+        // and add a start transition callback for each of them
+        const state = finishedWork.memoizedState;
+        // TODO Since it's a mutable field, this should live on the FiberRoot
+        if (state.transitions === null) {
+          state.transitions = new Set([]);
+        }
+        const pendingTransitions = state.transitions;
+        const pendingSuspenseBoundaries = state.pendingSuspenseBoundaries;
+
+        // Initial render
         if (committedTransitions !== null) {
           committedTransitions.forEach(transition => {
-            // TODO(luna) Do we want to log TransitionStart in the startTransition callback instead?
             addTransitionStartCallbackToPendingTransition({
               transitionName: transition.name,
               startTime: transition.startTime,
             });
+            pendingTransitions.add(transition);
+          });
 
-            addTransitionCompleteCallbackToPendingTransition({
-              transitionName: transition.name,
-              startTime: transition.startTime,
+          if (
+            pendingSuspenseBoundaries === null ||
+            pendingSuspenseBoundaries.size === 0
+          ) {
+            pendingTransitions.forEach(transition => {
+              addTransitionCompleteCallbackToPendingTransition({
+                transitionName: transition.name,
+                startTime: transition.startTime,
+              });
             });
-          });
+          }
 
           clearTransitionsForLanes(finishedRoot, committedLanes);
-          finishedWork.memoizedState.transitions = null;
+        }
+
+        // If there are no more pending suspense boundaries we
+        // clear the transitions because they are all complete.
+        if (
+          pendingSuspenseBoundaries === null ||
+          pendingSuspenseBoundaries.size === 0
+        ) {
+          state.transitions = null;
         }
       }
       break;
@@ -2800,9 +2914,44 @@ function commitPassiveMountOnFiber(
       }
 
       if (enableTransitionTracing) {
-        // TODO: Add code to actually process the update queue
+        const isFallback = finishedWork.memoizedState;
+        const queue = (finishedWork.updateQueue: any);
+        const rootMemoizedState = finishedRoot.current.memoizedState;
+
+        if (queue !== null) {
+          // We have one instance of the pendingSuspenseBoundaries map.
+          // We only need one because we update it during the commit phase.
+          // We instantiate a new Map if we haven't already
+          if (rootMemoizedState.pendingSuspenseBoundaries === null) {
+            rootMemoizedState.pendingSuspenseBoundaries = new Map();
+          }
+
+          if (isFallback) {
+            const transitions = queue.transitions;
+            let prevTransitions = finishedWork.memoizedState.transitions;
+            // Add all the transitions saved in the update queue during
+            // the render phase (ie the transitions associated with this boundary)
+            // into the transitions set.
+            if (transitions !== null) {
+              if (prevTransitions === null) {
+                // We only have one instance of the transitions set
+                // because we update it only during the commit phase. We
+                // will create the set on a as needed basis in the commit phase
+                finishedWork.memoizedState.transitions = prevTransitions = new Set();
+              }
+
+              transitions.forEach(transition => {
+                prevTransitions.add(transition);
+              });
+            }
+          }
+        }
+
+        commitTransitionProgress(finishedRoot, finishedWork);
+
         finishedWork.updateQueue = null;
       }
+
       break;
     }
     case CacheComponent: {

filename: packages/react-reconciler/src/__tests__/ReactTransitionTracing-test.js
patch: @@ -13,9 +13,14 @@ let ReactNoop;
 let Scheduler;
 let act;
 
+let getCacheForType;
 let useState;
+let Suspense;
 let startTransition;
 
+let caches;
+let seededCache;
+
 describe('ReactInteractionTracing', () => {
   beforeEach(() => {
     jest.resetModules();
@@ -28,13 +33,121 @@ describe('ReactInteractionTracing', () => {
 
     useState = React.useState;
     startTransition = React.startTransition;
+    Suspense = React.Suspense;
+
+    getCacheForType = React.unstable_getCacheForType;
+
+    caches = [];
+    seededCache = null;
   });
 
+  function createTextCache() {
+    if (seededCache !== null) {
+      const cache = seededCache;
+      seededCache = null;
+      return cache;
+    }
+
+    const data = new Map();
+    const cache = {
+      data,
+      resolve(text) {
+        const record = data.get(text);
+
+        if (record === undefined) {
+          const newRecord = {
+            status: 'resolved',
+            value: text,
+          };
+          data.set(text, newRecord);
+        } else if (record.status === 'pending') {
+          const thenable = record.value;
+          record.status = 'resolved';
+          record.value = text;
+          thenable.pings.forEach(t => t());
+        }
+      },
+      reject(text, error) {
+        const record = data.get(text);
+        if (record === undefined) {
+          const newRecord = {
+            status: 'rejected',
+            value: error,
+          };
+          data.set(text, newRecord);
+        } else if (record.status === 'pending') {
+          const thenable = record.value;
+          record.status = 'rejected';
+          record.value = error;
+          thenable.pings.forEach(t => t());
+        }
+      },
+    };
+    caches.push(cache);
+    return cache;
+  }
+
+  function readText(text) {
+    const textCache = getCacheForType(createTextCache);
+    const record = textCache.data.get(text);
+    if (record !== undefined) {
+      switch (record.status) {
+        case 'pending':
+          Scheduler.unstable_yieldValue(`Suspend [${text}]`);
+          throw record.value;
+        case 'rejected':
+          Scheduler.unstable_yieldValue(`Error [${text}]`);
+          throw record.value;
+        case 'resolved':
+          return record.value;
+      }
+    } else {
+      Scheduler.unstable_yieldValue(`Suspend [${text}]`);
+
+      const thenable = {
+        pings: [],
+        then(resolve) {
+          if (newRecord.status === 'pending') {
+            thenable.pings.push(resolve);
+          } else {
+            Promise.resolve().then(() => resolve(newRecord.value));
+          }
+        },
+      };
+
+      const newRecord = {
+        status: 'pending',
+        value: thenable,
+      };
+      textCache.data.set(text, newRecord);
+
+      throw thenable;
+    }
+  }
+
+  function AsyncText({text}) {
+    const fullText = readText(text);
+    Scheduler.unstable_yieldValue(fullText);
+    return fullText;
+  }
+
   function Text({text}) {
     Scheduler.unstable_yieldValue(text);
     return text;
   }
 
+  function resolveMostRecentTextCache(text) {
+    if (caches.length === 0) {
+      throw Error('Cache does not exist');
+    } else {
+      // Resolve the most recently created cache. An older cache can by
+      // resolved with `caches[index].resolve(text)`.
+      caches[caches.length - 1].resolve(text);
+    }
+  }
+
+  const resolveText = resolveMostRecentTextCache;
+
   function advanceTimers(ms) {
     // Note: This advances Jest's virtual time but not React's. Use
     // ReactNoop.expire for that.
@@ -98,4 +211,72 @@ describe('ReactInteractionTracing', () => {
       });
     });
   });
+
+  // @gate enableTransitionTracing
+  it('should correctly trace interactions for async roots', async () => {
+    const transitionCallbacks = {
+      onTransitionStart: (name, startTime) => {
+        Scheduler.unstable_yieldValue(
+          `onTransitionStart(${name}, ${startTime})`,
+        );
+      },
+      onTransitionComplete: (name, startTime, endTime) => {
+        Scheduler.unstable_yieldValue(
+          `onTransitionComplete(${name}, ${startTime}, ${endTime})`,
+        );
+      },
+    };
+    let navigateToPageTwo;
+    function App() {
+      const [navigate, setNavigate] = useState(false);
+      navigateToPageTwo = () => {
+        setNavigate(true);
+      };
+
+      return (
+        <div>
+          {navigate ? (
+            <Suspense
+              fallback={<Text text="Loading..." />}
+              name="suspense page">
+              <AsyncText text="Page Two" />
+            </Suspense>
+          ) : (
+            <Text text="Page One" />
+          )}
+        </div>
+      );
+    }
+
+    const root = ReactNoop.createRoot({transitionCallbacks});
+    await act(async () => {
+      root.render(<App />);
+      ReactNoop.expire(1000);
+      await advanceTimers(1000);
+
+      expect(Scheduler).toFlushAndYield(['Page One']);
+    });
+
+    await act(async () => {
+      startTransition(() => navigateToPageTwo(), {name: 'page transition'});
+
+      ReactNoop.expire(1000);
+      await advanceTimers(1000);
+
+      expect(Scheduler).toFlushAndYield([
+        'Suspend [Page Two]',
+        'Loading...',
+        'onTransitionStart(page transition, 1000)',
+      ]);
+
+      ReactNoop.expire(1000);
+      await advanceTimers(1000);
+      await resolveText('Page Two');
+
+      expect(Scheduler).toFlushAndYield([
+        'Page Two',
+        'onTransitionComplete(page transition, 1000, 3000)',
+      ]);
+    });
+  });
 });