Added several new interaction tracing tests and edge case fixes

Author: Brian Vaughn

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -76,6 +76,7 @@ import {
 } from './ReactFiberHostConfig';
 import {
   captureCommitPhaseError,
+  markSuspenseDidTimeout,
   requestCurrentTime,
   scheduleWork,
 } from './ReactFiberScheduler';
@@ -361,6 +362,12 @@ function commitLifeCycles(
           // entire queue. Any non-null value works.
           // $FlowFixMe - Intentionally using a value other than an UpdateQueue.
           finishedWork.updateQueue = emptyObject;
+          if (enableSchedulerTracing) {
+            // Handles the special case of a suspended sync render.
+            // In this case we don't decrement the interaction count on Placeholder commit.
+            // Instead we wait until the suspended promise has resolved.
+            markSuspenseDidTimeout();
+          }
           scheduleWork(finishedWork, Sync);
         } else {
           // In strict mode, the Update effect is used to record the time at

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -260,8 +260,14 @@ 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.
+// Instead we wait until the suspended promise has resolved.
 let suspenseDidTimeout: boolean = false;
 
+// Do not remove the current interactions from the priority map on commit.
+// This covers a special case of sync rendering with suspense.
+// In this case we leave interactions in the Map until the suspended promise resolves.
+let preservePendingInteractionsOnCommit: boolean = true;
+
 let stashedWorkInProgressProperties;
 let replayUnitOfWork;
 let isReplayingFailedUnitOfWork;
@@ -583,11 +589,17 @@ function commitRoot(root: FiberRoot, finishedWork: Fiber): void {
     // 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(
+    const pendingInteractionMap = root.pendingInteractionMap;
+    pendingInteractionMap.forEach(
       (scheduledInteractions, scheduledExpirationTime) => {
         if (scheduledExpirationTime <= committedExpirationTime) {
           committedInteractions.push(...Array.from(scheduledInteractions));
-          root.pendingInteractionMap.delete(scheduledExpirationTime);
+
+          // Don't delete interactions from the map if we're frozen due to suspense.
+          // In this case, the interactions should be moved to the suspense retry time.
+          if (!suspenseDidTimeout && !preservePendingInteractionsOnCommit) {
+            pendingInteractionMap.delete(scheduledExpirationTime);
+          }
         }
       },
     );
@@ -1174,14 +1186,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 +1205,9 @@ function renderRoot(
     root.pendingCommitExpirationTime = NoWork;
 
     if (enableSchedulerTracing) {
+      preservePendingInteractionsOnCommit = suspenseDidTimeout;
+      suspenseDidTimeout = 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 +1251,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 +1418,7 @@ function renderRoot(
 
   if (enableSuspense && !isExpired && nextLatestAbsoluteTimeoutMs !== -1) {
     // The tree was suspended.
+    suspenseDidTimeout = true;
     const suspendedExpirationTime = expirationTime;
     markSuspendedPriorityLevel(root, suspendedExpirationTime);
 
@@ -1603,29 +1619,38 @@ 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
       const currentTime = requestCurrentTime();
       retryTime = computeExpirationForFiber(currentTime, fiber);
       markPendingPriorityLevel(root, retryTime);
+
+      if (enableSchedulerTracing) {
+        // Interaction counts should not get decremented until suspended work resolves.
+        // Update their position in the map to account for this expiration time change.
+        // This ensures they aren't decremented prematurely (e.g. by high pri interleaved work).
+        const pendingInteractionMap = root.pendingInteractionMap;
+        const suspendedInteractions = pendingInteractionMap.get(suspendedTime);
+
+        if (suspendedInteractions != null) {
+          const retryInteractions = pendingInteractionMap.get(retryTime);
+          if (retryInteractions != null) {
+            suspendedInteractions.forEach(interaction => {
+              retryInteractions.add(interaction);
+            });
+          } else {
+            pendingInteractionMap.set(retryTime, suspendedInteractions);
+          }
+        }
+      }
     }
 
     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);
     }
   }
 }
@@ -1680,49 +1705,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();
 
@@ -1745,7 +1727,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 (
@@ -1940,7 +1952,6 @@ function onTimeout(root, finishedWork, suspendedExpirationTime) {
       // Because we know we still need to do more work in this case.
       suspenseDidTimeout = true;
       flushRoot(root, suspendedExpirationTime);
-      suspenseDidTimeout = false;
     } else {
       flushRoot(root, suspendedExpirationTime);
     }
@@ -2511,6 +2522,12 @@ function flushControlled(fn: () => mixed): void {
   }
 }
 
+function markSuspenseDidTimeout() {
+  if (enableSchedulerTracing) {
+    suspenseDidTimeout = true;
+  }
+}
+
 export {
   requestCurrentTime,
   computeExpirationForFiber,
@@ -2533,4 +2550,5 @@ export {
   interactiveUpdates,
   flushInteractiveUpdates,
   computeUniqueAsyncExpiration,
+  markSuspenseDidTimeout,
 };