Refactored interaction tracing to fix ref counting bug

Author: Brian Vaughn

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -41,13 +41,15 @@ import {
   Update,
   Ref,
 } from 'shared/ReactSideEffectTags';
+import {captureWillSyncRenderPlaceholder} from './ReactFiberScheduler';
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 import {
   enableGetDerivedStateFromCatch,
   enableSuspense,
   debugRenderPhaseSideEffects,
   debugRenderPhaseSideEffectsForStrictMode,
   enableProfilerTimer,
+  enableSchedulerTracing,
 } from 'shared/ReactFeatureFlags';
 import invariant from 'shared/invariant';
 import getComponentName from 'shared/getComponentName';
@@ -825,6 +827,13 @@ function updatePlaceholderComponent(
 
     let nextDidTimeout;
     if (current !== null && workInProgress.updateQueue !== null) {
+      if (enableSchedulerTracing) {
+        // Handle special case of rendering a Placeholder for a sync, suspended tree.
+        // We flag this to properly trace and count interactions.
+        // Otherwise interaction pending count will be decremented too many times.
+        captureWillSyncRenderPlaceholder();
+      }
+
       // We're outside strict mode. Something inside this Placeholder boundary
       // suspended during the last commit. Switch to the placholder.
       workInProgress.updateQueue = null;

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -260,7 +260,8 @@ let interruptedBy: Fiber | null = null;
 
 // Do not decrement interaction counts in the event of suspense timeouts.
 // This would lead to prematurely calling the interaction-complete hook.
-let suspenseDidTimeout: boolean = false;
+// Instead we wait until the suspended promise has resolved.
+let interactionsHaveBeenSuspended: boolean = false;
 
 let stashedWorkInProgressProperties;
 let replayUnitOfWork;
@@ -566,31 +567,11 @@ function commitRoot(root: FiberRoot, finishedWork: Fiber): void {
   markCommittedPriorityLevels(root, earliestRemainingTimeBeforeCommit);
 
   let prevInteractions: Set<Interaction> = (null: any);
-  let committedInteractions: Array<Interaction> = enableSchedulerTracing
-    ? []
-    : (null: any);
   if (enableSchedulerTracing) {
     // Restore any pending interactions at this point,
     // So that cascading work triggered during the render phase will be accounted for.
     prevInteractions = __interactionsRef.current;
     __interactionsRef.current = root.memoizedInteractions;
-
-    // We are potentially finished with the current batch of interactions.
-    // So we should clear them out of the pending interaction map.
-    // We do this at the start of commit in case cascading work is scheduled by commit phase lifecycles.
-    // In that event, interaction data may be added back into the pending map for a future commit.
-    // We also store the interactions we are about to commit so that we can notify subscribers after we're done.
-    // These are stored as an Array rather than a Set,
-    // Because the same interaction may be pending for multiple expiration times,
-    // In which case it's important that we decrement the count the right number of times after finishing.
-    root.pendingInteractionMap.forEach(
-      (scheduledInteractions, scheduledExpirationTime) => {
-        if (scheduledExpirationTime <= committedExpirationTime) {
-          committedInteractions.push(...Array.from(scheduledInteractions));
-          root.pendingInteractionMap.delete(scheduledExpirationTime);
-        }
-      },
-    );
   }
 
   // Reset this to null before calling lifecycles
@@ -789,27 +770,38 @@ function commitRoot(root: FiberRoot, finishedWork: Fiber): void {
         unhandledError = error;
       }
     } finally {
-      // Don't update interaction counts if we're frozen due to suspense.
-      // In this case, we can skip the completed-work check entirely.
-      if (!suspenseDidTimeout) {
-        // Now that we're done, check the completed batch of interactions.
-        // If no more work is outstanding for a given interaction,
-        // We need to notify the subscribers that it's finished.
-        committedInteractions.forEach(interaction => {
-          interaction.__count--;
-          if (subscriber !== null && interaction.__count === 0) {
-            try {
-              subscriber.onInteractionScheduledWorkCompleted(interaction);
-            } catch (error) {
-              // It's not safe for commitRoot() to throw.
-              // Store the error for now and we'll re-throw in finishRendering().
-              if (!hasUnhandledError) {
-                hasUnhandledError = true;
-                unhandledError = error;
-              }
+      if (!interactionsHaveBeenSuspended) {
+        // Clear completed interactions from the pending Map.
+        // Unless the render was suspended or cascading work was scheduled,
+        // In which case– leave pending interactions until the subsequent render.
+        const pendingInteractionMap = root.pendingInteractionMap;
+        pendingInteractionMap.forEach(
+          (scheduledInteractions, scheduledExpirationTime) => {
+            if (
+              earliestRemainingTimeAfterCommit === NoWork ||
+              scheduledExpirationTime < earliestRemainingTimeAfterCommit
+            ) {
+              pendingInteractionMap.delete(scheduledExpirationTime);
+
+              scheduledInteractions.forEach(interaction => {
+                interaction.__count--;
+
+                if (subscriber !== null && interaction.__count === 0) {
+                  try {
+                    subscriber.onInteractionScheduledWorkCompleted(interaction);
+                  } catch (error) {
+                    // It's not safe for commitRoot() to throw.
+                    // Store the error for now and we'll re-throw in finishRendering().
+                    if (!hasUnhandledError) {
+                      hasUnhandledError = true;
+                      unhandledError = error;
+                    }
+                  }
+                }
+              });
             }
-          }
-        });
+          },
+        );
       }
     }
   }
@@ -1174,14 +1166,6 @@ function renderRoot(
 
   const expirationTime = root.nextExpirationTimeToWorkOn;
 
-  let prevInteractions: Set<Interaction> = (null: any);
-  if (enableSchedulerTracing) {
-    // We're about to start new traced work.
-    // Restore pending interactions so cascading work triggered during the render phase will be accounted for.
-    prevInteractions = __interactionsRef.current;
-    __interactionsRef.current = root.memoizedInteractions;
-  }
-
   // Check if we're starting from a fresh stack, or if we're resuming from
   // previously yielded work.
   if (
@@ -1201,6 +1185,11 @@ function renderRoot(
     root.pendingCommitExpirationTime = NoWork;
 
     if (enableSchedulerTracing) {
+      // Reset this flag once we start rendering a new root or at a new priority.
+      // This might indicate that suspended work has completed.
+      // If not, the flag will be reset.
+      interactionsHaveBeenSuspended = false;
+
       // 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.
@@ -1244,6 +1233,14 @@ function renderRoot(
     }
   }
 
+  let prevInteractions: Set<Interaction> = (null: any);
+  if (enableSchedulerTracing) {
+    // We're about to start new traced work.
+    // Restore pending interactions so cascading work triggered during the render phase will be accounted for.
+    prevInteractions = __interactionsRef.current;
+    __interactionsRef.current = root.memoizedInteractions;
+  }
+
   let didFatal = false;
 
   startWorkLoopTimer(nextUnitOfWork);
@@ -1403,6 +1400,9 @@ function renderRoot(
 
   if (enableSuspense && !isExpired && nextLatestAbsoluteTimeoutMs !== -1) {
     // The tree was suspended.
+    if (enableSchedulerTracing) {
+      interactionsHaveBeenSuspended = true;
+    }
     const suspendedExpirationTime = expirationTime;
     markSuspendedPriorityLevel(root, suspendedExpirationTime);
 
@@ -1600,6 +1600,7 @@ function retrySuspendedRoot(
     if (isPriorityLevelSuspended(root, suspendedTime)) {
       // Ping at the original level
       retryTime = suspendedTime;
+
       markPingedPriorityLevel(root, retryTime);
     } else {
       // Placeholder already timed out. Compute a new expiration time
@@ -1611,18 +1612,7 @@ function retrySuspendedRoot(
     scheduleWorkToRoot(fiber, retryTime);
     const rootExpirationTime = root.expirationTime;
     if (rootExpirationTime !== NoWork) {
-      if (enableSchedulerTracing) {
-        // Restore previous interactions so that new work is associated with them.
-        let prevInteractions = __interactionsRef.current;
-        __interactionsRef.current = root.memoizedInteractions;
-        // Because suspense timeouts do not decrement the interaction count,
-        // Continued suspense work should also not increment the count.
-        storeInteractionsForExpirationTime(root, rootExpirationTime, false);
-        requestWork(root, rootExpirationTime);
-        __interactionsRef.current = prevInteractions;
-      } else {
-        requestWork(root, rootExpirationTime);
-      }
+      requestWork(root, rootExpirationTime);
     }
   }
 }
@@ -1677,49 +1667,6 @@ function scheduleWorkToRoot(fiber: Fiber, expirationTime): FiberRoot | null {
   return null;
 }
 
-function storeInteractionsForExpirationTime(
-  root: FiberRoot,
-  expirationTime: ExpirationTime,
-  updateInteractionCounts: boolean,
-): void {
-  if (!enableSchedulerTracing) {
-    return;
-  }
-
-  const interactions = __interactionsRef.current;
-  if (interactions.size > 0) {
-    const pendingInteractions = root.pendingInteractionMap.get(expirationTime);
-    if (pendingInteractions != null) {
-      interactions.forEach(interaction => {
-        if (updateInteractionCounts && !pendingInteractions.has(interaction)) {
-          // Update the pending async work count for previously unscheduled interaction.
-          interaction.__count++;
-        }
-
-        pendingInteractions.add(interaction);
-      });
-    } else {
-      root.pendingInteractionMap.set(expirationTime, new Set(interactions));
-
-      // Update the pending async work count for the current interactions.
-      if (updateInteractionCounts) {
-        interactions.forEach(interaction => {
-          interaction.__count++;
-        });
-      }
-    }
-
-    const subscriber = __subscriberRef.current;
-    if (subscriber !== null) {
-      const threadID = computeThreadID(
-        expirationTime,
-        root.interactionThreadID,
-      );
-      subscriber.onWorkScheduled(interactions, threadID);
-    }
-  }
-}
-
 function scheduleWork(fiber: Fiber, expirationTime: ExpirationTime) {
   recordScheduleUpdate();
 
@@ -1742,7 +1689,37 @@ function scheduleWork(fiber: Fiber, expirationTime: ExpirationTime) {
   }
 
   if (enableSchedulerTracing) {
-    storeInteractionsForExpirationTime(root, expirationTime, true);
+    const interactions = __interactionsRef.current;
+    if (interactions.size > 0) {
+      const pendingInteractionMap = root.pendingInteractionMap;
+      const pendingInteractions = pendingInteractionMap.get(expirationTime);
+      if (pendingInteractions != null) {
+        interactions.forEach(interaction => {
+          if (!pendingInteractions.has(interaction)) {
+            // Update the pending async work count for previously unscheduled interaction.
+            interaction.__count++;
+          }
+
+          pendingInteractions.add(interaction);
+        });
+      } else {
+        pendingInteractionMap.set(expirationTime, new Set(interactions));
+
+        // Update the pending async work count for the current interactions.
+        interactions.forEach(interaction => {
+          interaction.__count++;
+        });
+      }
+
+      const subscriber = __subscriberRef.current;
+      if (subscriber !== null) {
+        const threadID = computeThreadID(
+          expirationTime,
+          root.interactionThreadID,
+        );
+        subscriber.onWorkScheduled(interactions, threadID);
+      }
+    }
   }
 
   if (
@@ -1935,9 +1912,8 @@ function onTimeout(root, finishedWork, suspendedExpirationTime) {
     if (enableSchedulerTracing) {
       // Don't update pending interaction counts for suspense timeouts,
       // Because we know we still need to do more work in this case.
-      suspenseDidTimeout = true;
+      interactionsHaveBeenSuspended = true;
       flushRoot(root, suspendedExpirationTime);
-      suspenseDidTimeout = false;
     } else {
       flushRoot(root, suspendedExpirationTime);
     }
@@ -2508,6 +2484,12 @@ function flushControlled(fn: () => mixed): void {
   }
 }
 
+function captureWillSyncRenderPlaceholder() {
+  if (enableSchedulerTracing) {
+    interactionsHaveBeenSuspended = true;
+  }
+}
+
 export {
   requestCurrentTime,
   computeExpirationForFiber,
@@ -2530,4 +2512,5 @@ export {
   interactiveUpdates,
   flushInteractiveUpdates,
   computeUniqueAsyncExpiration,
+  captureWillSyncRenderPlaceholder,
 };

packages/react-reconciler/src/ReactFiberUnwindWork.js

@@ -36,6 +36,7 @@ import {
 import {
   enableGetDerivedStateFromCatch,
   enableSuspense,
+  enableSchedulerTracing,
 } from 'shared/ReactFeatureFlags';
 import {StrictMode, AsyncMode} from './ReactTypeOfMode';
 
@@ -60,6 +61,7 @@ import {
   markLegacyErrorBoundaryAsFailed,
   isAlreadyFailedLegacyErrorBoundary,
   retrySuspendedRoot,
+  captureWillSyncRenderPlaceholder,
 } from './ReactFiberScheduler';
 import {Sync} from './ReactFiberExpirationTime';
 
@@ -236,6 +238,13 @@ function throwException(
           if ((workInProgress.mode & StrictMode) === NoEffect) {
             workInProgress.effectTag |= UpdateEffect;
 
+            if (enableSchedulerTracing) {
+              // Handles the special case of unwinding a suspended sync render.
+              // We flag this to properly trace and count interactions.
+              // Otherwise interaction pending count will be decremented too many times.
+              captureWillSyncRenderPlaceholder();
+            }
+
             // Unmount the source fiber's children
             const nextChildren = null;
             reconcileChildren(