Track which fibers scheduled the current render work

Author: Brian Vaughn

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -28,6 +28,7 @@ import {
   enableSchedulerTracing,
   enableProfilerTimer,
   enableSuspenseServerRenderer,
+  enableUpdaterTracking,
   enableEventAPI,
 } from 'shared/ReactFeatureFlags';
 import {
@@ -104,6 +105,7 @@ import {
   captureCommitPhaseError,
   requestCurrentTime,
   resolveRetryThenable,
+  restorePendingUpdaters,
 } from './ReactFiberScheduler';
 import {
   NoEffect as NoHookEffect,
@@ -1169,7 +1171,12 @@ function commitDeletion(current: Fiber): void {
   detachFiber(current);
 }
 
-function commitWork(current: Fiber | null, finishedWork: Fiber): void {
+function commitWork(
+  finishedRoot: FiberRoot,
+  current: Fiber | null,
+  finishedWork: Fiber,
+  committedExpirationTime: ExpirationTime,
+): void {
   if (!supportsMutation) {
     switch (finishedWork.tag) {
       case FunctionComponent:
@@ -1185,7 +1192,11 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
         return;
       }
       case SuspenseComponent: {
-        commitSuspenseComponent(finishedWork);
+        commitSuspenseComponent(
+          finishedRoot,
+          finishedWork,
+          committedExpirationTime,
+        );
         return;
       }
     }
@@ -1259,7 +1270,11 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
       return;
     }
     case SuspenseComponent: {
-      commitSuspenseComponent(finishedWork);
+      commitSuspenseComponent(
+        finishedRoot,
+        finishedWork,
+        committedExpirationTime,
+      );
       return;
     }
     case IncompleteClassComponent: {
@@ -1278,7 +1293,11 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
   }
 }
 
-function commitSuspenseComponent(finishedWork: Fiber) {
+function commitSuspenseComponent(
+  finishedRoot: FiberRoot,
+  finishedWork: Fiber,
+  committedExpirationTime: ExpirationTime,
+) {
   let newState: SuspenseState | null = finishedWork.memoizedState;
 
   let newDidTimeout;
@@ -1321,6 +1340,10 @@ function commitSuspenseComponent(finishedWork: Fiber) {
         if (enableSchedulerTracing) {
           retry = Schedule_tracing_wrap(retry);
         }
+        if (enableUpdaterTracking) {
+          // If we have pending work still, restore the original updaters
+          restorePendingUpdaters(finishedRoot, committedExpirationTime);
+        }
         retryCache.add(thenable);
         thenable.then(retry, retry);
       }

packages/react-reconciler/src/ReactFiberRoot.js

@@ -17,7 +17,10 @@ import type {Interaction} from 'scheduler/src/Tracing';
 import {noTimeout} from './ReactFiberHostConfig';
 import {createHostRootFiber} from './ReactFiber';
 import {NoWork} from './ReactFiberExpirationTime';
-import {enableSchedulerTracing} from 'shared/ReactFeatureFlags';
+import {
+  enableSchedulerTracing,
+  enableUpdaterTracking,
+} from 'shared/ReactFeatureFlags';
 import {unstable_getThreadID} from 'scheduler/tracing';
 
 // TODO: This should be lifted into the renderer.
@@ -30,6 +33,9 @@ export type Batch = {
 
 export type PendingInteractionMap = Map<ExpirationTime, Set<Interaction>>;
 
+// Map of expiration time to all pending "updaters" which in turn is a map of Fibers to reference counts.
+export type PendingUpdatersMap = Map<ExpirationTime, Set<Fiber>>;
+
 type BaseFiberRootProperties = {|
   // The type of root (legacy, batched, concurrent, etc.)
   tag: RootTag,
@@ -83,6 +89,13 @@ type ProfilingOnlyFiberRootProperties = {|
   pendingInteractionMap: PendingInteractionMap,
 |};
 
+// The following attributes are only used by DevTools and are only present in DEV builds.
+// They enable DevTools Profiler UI to show which Fiber(s) scheduled a given commit.
+type UpdaterTrackingOnlyFiberRootProperties = {|
+  memoizedUpdaters: Set<Fiber>,
+  pendingUpdatersMap: PendingUpdatersMap,
+|};
+
 // Exported FiberRoot type includes all properties,
 // To avoid requiring potentially error-prone :any casts throughout the project.
 // Profiling properties are only safe to access in profiling builds (when enableSchedulerTracing is true).
@@ -91,6 +104,7 @@ type ProfilingOnlyFiberRootProperties = {|
 export type FiberRoot = {
   ...BaseFiberRootProperties,
   ...ProfilingOnlyFiberRootProperties,
+  ...UpdaterTrackingOnlyFiberRootProperties,
 };
 
 function FiberRootNode(containerInfo, tag, hydrate) {
@@ -117,6 +131,11 @@ function FiberRootNode(containerInfo, tag, hydrate) {
     this.memoizedInteractions = new Set();
     this.pendingInteractionMap = new Map();
   }
+
+  if (enableUpdaterTracking) {
+    this.memoizedUpdaters = new Set();
+    this.pendingUpdatersMap = new Map();
+  }
 }
 
 export function createFiberRoot(

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -19,6 +19,7 @@ import type {Interaction} from 'scheduler/src/Tracing';
 import {
   warnAboutDeprecatedLifecycles,
   enableUserTimingAPI,
+  enableUpdaterTracking,
   enableSuspenseServerRenderer,
   replayFailedUnitOfWorkWithInvokeGuardedCallback,
   enableProfilerTimer,
@@ -337,6 +338,16 @@ export function scheduleUpdateOnFiber(
     return;
   }
 
+  if (enableUpdaterTracking) {
+    const pendingUpdatersMap = root.pendingUpdatersMap;
+    let updaters = pendingUpdatersMap.get(expirationTime);
+    if (updaters == null) {
+      updaters = new Set();
+      pendingUpdatersMap.set(expirationTime, updaters);
+    }
+    updaters.add(fiber);
+  }
+
   root.pingTime = NoWork;
 
   checkForInterruption(fiber, expirationTime);
@@ -1292,6 +1303,12 @@ function commitRootImpl(root) {
     // This usually means we've finished all the work, but it can also happen
     // when something gets downprioritized during render, like a hidden tree.
     root.lastPendingTime = firstPendingTimeBeforeCommit;
+
+    if (enableUpdaterTracking) {
+      if (firstPendingTimeBeforeCommit !== NoWork) {
+        restorePendingUpdaters(root, root.lastPendingTime);
+      }
+    }
   }
 
   if (root === workInProgressRoot) {
@@ -1377,7 +1394,13 @@ function commitRootImpl(root) {
     nextEffect = firstEffect;
     do {
       if (__DEV__) {
-        invokeGuardedCallback(null, commitMutationEffects, null);
+        invokeGuardedCallback(
+          null,
+          commitMutationEffects,
+          null,
+          root,
+          expirationTime,
+        );
         if (hasCaughtError()) {
           invariant(nextEffect !== null, 'Should be working on an effect.');
           const error = clearCaughtError();
@@ -1386,7 +1409,7 @@ function commitRootImpl(root) {
         }
       } else {
         try {
-          commitMutationEffects();
+          commitMutationEffects(root, expirationTime);
         } catch (error) {
           invariant(nextEffect !== null, 'Should be working on an effect.');
           captureCommitPhaseError(nextEffect, error);
@@ -1540,7 +1563,10 @@ function commitBeforeMutationEffects() {
   }
 }
 
-function commitMutationEffects() {
+function commitMutationEffects(
+  root: FiberRoot,
+  committedExpirationTime: ExpirationTime,
+) {
   // TODO: Should probably move the bulk of this function to commitWork.
   while (nextEffect !== null) {
     setCurrentDebugFiberInDEV(nextEffect);
@@ -1582,12 +1608,12 @@ function commitMutationEffects() {
 
         // Update
         const current = nextEffect.alternate;
-        commitWork(current, nextEffect);
+        commitWork(root, current, nextEffect, committedExpirationTime);
         break;
       }
       case Update: {
         const current = nextEffect.alternate;
-        commitWork(current, nextEffect);
+        commitWork(root, current, nextEffect, committedExpirationTime);
         break;
       }
       case Deletion: {
@@ -2161,6 +2187,24 @@ function warnIfNotCurrentlyActingUpdatesInDEV(fiber: Fiber): void {
   }
 }
 
+export function restorePendingUpdaters(
+  root: FiberRoot,
+  expirationTime: ExpirationTime,
+): void {
+  if (!enableUpdaterTracking) {
+    return;
+  }
+  const pendingUpdatersMap = root.pendingUpdatersMap;
+  let updaters = pendingUpdatersMap.get(expirationTime);
+  if (updaters == null) {
+    updaters = new Set();
+    pendingUpdatersMap.set(expirationTime, updaters);
+  }
+  root.memoizedUpdaters.forEach(schedulingFiber => {
+    ((updaters: any): Set<Fiber>).add(schedulingFiber);
+  });
+}
+
 export const warnIfNotCurrentlyActingUpdatesInDev = warnIfNotCurrentlyActingUpdatesInDEV;
 
 let componentsWithSuspendedDiscreteUpdates = null;
@@ -2277,42 +2321,58 @@ function schedulePendingInteraction(root, expirationTime) {
 
 function startWorkOnPendingInteraction(root, expirationTime) {
   // This is called when new work is started on a root.
-  if (!enableSchedulerTracing) {
-    return;
-  }
 
-  // Determine which interactions this batch of work currently includes, So that
-  // we can accurately attribute time spent working on it, And so that cascading
-  // work triggered during the render phase will be associated with it.
-  const interactions: Set<Interaction> = new Set();
-  root.pendingInteractionMap.forEach(
-    (scheduledInteractions, scheduledExpirationTime) => {
+  if (enableUpdaterTracking) {
+    const memoizedUpdaters: Set<Fiber> = new Set();
+    const pendingUpdatersMap = root.pendingUpdatersMap;
+    pendingUpdatersMap.forEach((updaters, scheduledExpirationTime) => {
       if (scheduledExpirationTime >= expirationTime) {
-        scheduledInteractions.forEach(interaction =>
-          interactions.add(interaction),
-        );
+        pendingUpdatersMap.delete(scheduledExpirationTime);
+        updaters.forEach(fiber => memoizedUpdaters.add(fiber));
       }
-    },
-  );
+    });
 
-  // Store the current set of interactions on the FiberRoot for a few reasons:
-  // We can re-use it in hot functions like renderRoot() without having to
-  // recalculate it. We will also use it in commitWork() to pass to any Profiler
-  // onRender() hooks. This also provides DevTools with a way to access it when
-  // the onCommitRoot() hook is called.
-  root.memoizedInteractions = interactions;
+    // Store the current set of interactions on the FiberRoot for a few reasons:
+    // We can re-use it in hot functions like renderRoot() without having to
+    // recalculate it. This also provides DevTools with a way to access it when
+    // the onCommitRoot() hook is called.
+    root.memoizedUpdaters = memoizedUpdaters;
+  }
 
-  if (interactions.size > 0) {
-    const subscriber = __subscriberRef.current;
-    if (subscriber !== null) {
-      const threadID = computeThreadID(root, expirationTime);
-      try {
-        subscriber.onWorkStarted(interactions, threadID);
-      } catch (error) {
-        // If the subscriber throws, rethrow it in a separate task
-        scheduleCallback(ImmediatePriority, () => {
-          throw error;
-        });
+  if (enableSchedulerTracing) {
+    // Determine which interactions this batch of work currently includes, So that
+    // we can accurately attribute time spent working on it, And so that cascading
+    // work triggered during the render phase will be associated with it.
+    const interactions: Set<Interaction> = new Set();
+    root.pendingInteractionMap.forEach(
+      (scheduledInteractions, scheduledExpirationTime) => {
+        if (scheduledExpirationTime >= expirationTime) {
+          scheduledInteractions.forEach(interaction =>
+            interactions.add(interaction),
+          );
+        }
+      },
+    );
+
+    // Store the current set of interactions on the FiberRoot for a few reasons:
+    // We can re-use it in hot functions like renderRoot() without having to
+    // recalculate it. We will also use it in commitWork() to pass to any Profiler
+    // onRender() hooks. This also provides DevTools with a way to access it when
+    // the onCommitRoot() hook is called.
+    root.memoizedInteractions = interactions;
+
+    if (interactions.size > 0) {
+      const subscriber = __subscriberRef.current;
+      if (subscriber !== null) {
+        const threadID = computeThreadID(root, expirationTime);
+        try {
+          subscriber.onWorkStarted(interactions, threadID);
+        } catch (error) {
+          // If the subscriber throws, rethrow it in a separate task
+          scheduleCallback(ImmediatePriority, () => {
+            throw error;
+          });
+        }
       }
     }
   }

packages/react-reconciler/src/ReactFiberUnwindWork.js

@@ -40,6 +40,7 @@ import {
 import {
   enableSchedulerTracing,
   enableSuspenseServerRenderer,
+  enableUpdaterTracking,
   enableEventAPI,
 } from 'shared/ReactFeatureFlags';
 import {NoMode, BatchedMode} from './ReactTypeOfMode';
@@ -76,6 +77,7 @@ import {
   pingSuspendedRoot,
   resolveRetryThenable,
   checkForWrongSuspensePriorityInDEV,
+  restorePendingUpdaters,
 } from './ReactFiberScheduler';
 
 import invariant from 'shared/invariant';
@@ -187,6 +189,10 @@ function attachPingListener(
     if (enableSchedulerTracing) {
       ping = Schedule_tracing_wrap(ping);
     }
+    if (enableUpdaterTracking) {
+      // If we have pending work still, restore the original updaters
+      restorePendingUpdaters(root, renderExpirationTime);
+    }
     thenable.then(ping, ping);
   }
 }
@@ -203,6 +209,11 @@ function throwException(
   // Its effect list is no longer valid.
   sourceFiber.firstEffect = sourceFiber.lastEffect = null;
 
+  if (enableUpdaterTracking) {
+    // If we have pending work still, restore the original updaters
+    restorePendingUpdaters(root, renderExpirationTime);
+  }
+
   if (
     value !== null &&
     typeof value === 'object' &&