Added some basic tests and continue schedulers across Suspense bounda…

…ries
facebook · May 15, 2019 · daf9abf · daf9abf
1 parent 49b5450
commit daf9abf
Show file tree

Hide file tree

Showing 5 changed files with 289 additions and 35 deletions.
diff --git a/packages/react-reconciler/src/ReactFiberCommitWork.js b/packages/react-reconciler/src/ReactFiberCommitWork.js
@@ -104,6 +104,7 @@ import {
   captureCommitPhaseError,
   requestCurrentTime,
   resolveRetryThenable,
+  restorePendingSchedulers,
 } from './ReactFiberScheduler';
 import {
   NoEffect as NoHookEffect,
@@ -1169,7 +1170,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 +1191,11 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
         return;
       }
       case SuspenseComponent: {
-        commitSuspenseComponent(finishedWork);
+        commitSuspenseComponent(
+          finishedRoot,
+          finishedWork,
+          committedExpirationTime,
+        );
         return;
       }
     }
@@ -1259,7 +1269,11 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
       return;
     }
     case SuspenseComponent: {
-      commitSuspenseComponent(finishedWork);
+      commitSuspenseComponent(
+        finishedRoot,
+        finishedWork,
+        committedExpirationTime,
+      );
       return;
     }
     case IncompleteClassComponent: {
@@ -1278,7 +1292,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;
@@ -1320,6 +1338,9 @@ function commitSuspenseComponent(finishedWork: Fiber) {
       if (!retryCache.has(thenable)) {
         if (enableSchedulerTracing) {
           retry = Schedule_tracing_wrap(retry);
+
+          // If we have pending work still, restore the original schedulers
+          restorePendingSchedulers(finishedRoot, committedExpirationTime);
         }
         retryCache.add(thenable);
         thenable.then(retry, retry);

diff --git a/packages/react-reconciler/src/ReactFiberRoot.js b/packages/react-reconciler/src/ReactFiberRoot.js
@@ -29,7 +29,9 @@ export type Batch = {
 };
 
 export type PendingInteractionMap = Map<ExpirationTime, Set<Interaction>>;
-export type PendingSchedulerMap = Map<ExpirationTime, Set<Fiber>>;
+
+// Map of expiration time to all pending "schedulers" which in turn is a map of Fibers to reference counts.
+export type PendingSchedulersMap = Map<ExpirationTime, Set<Fiber>>;
 
 type BaseFiberRootProperties = {|
   // The type of root (legacy, batched, concurrent, etc.)
@@ -86,7 +88,7 @@ type ProfilingOnlyFiberRootProperties = {|
   // Used to enable DevTools Profiler UI to show which Fibers scheduled a given commit.
   // May also be useful in the future to expose via the Profiler API somehow?
   memoizedSchedulers: Set<Fiber>,
-  pendingSchedulerMap: PendingSchedulerMap,
+  pendingSchedulersMap: PendingSchedulersMap,
 |};
 
 // Exported FiberRoot type includes all properties,
@@ -123,7 +125,7 @@ function FiberRootNode(containerInfo, tag, hydrate) {
     this.memoizedInteractions = new Set();
     this.pendingInteractionMap = new Map();
     this.memoizedSchedulers = new Set();
-    this.pendingSchedulerMap = new Map();
+    this.pendingSchedulersMap = new Map();
   }
 }
 

diff --git a/packages/react-reconciler/src/ReactFiberScheduler.js b/packages/react-reconciler/src/ReactFiberScheduler.js
@@ -337,13 +337,13 @@ export function scheduleUpdateOnFiber(
   }
 
   if (enableSchedulerTracing) {
-    const pendingSchedulerMap = root.pendingSchedulerMap;
-    const pendingSchedulers = pendingSchedulerMap.get(expirationTime);
-    if (pendingSchedulers != null) {
-      pendingSchedulers.add(fiber);
-    } else {
-      pendingSchedulerMap.set(expirationTime, new Set([fiber]));
+    const pendingSchedulersMap = root.pendingSchedulersMap;
+    let schedulers = pendingSchedulersMap.get(expirationTime);
+    if (schedulers == null) {
+      schedulers = new Set();
+      pendingSchedulersMap.set(expirationTime, schedulers);
     }
+    schedulers.add(fiber);
   }
 
   root.pingTime = NoWork;
@@ -1377,7 +1377,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 +1392,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 +1546,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 +1591,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: {
@@ -1843,7 +1852,12 @@ export function retryTimedOutBoundary(boundaryFiber: Fiber) {
   }
 }
 
-export function resolveRetryThenable(boundaryFiber: Fiber, thenable: Thenable) {
+export function resolveRetryThenable(
+  boundaryFiber: Fiber,
+  thenable: Thenable,
+  root: FiberRoot,
+  committedExpirationTime: ExpirationTime,
+) {
   let retryCache: WeakSet<Thenable> | Set<Thenable> | null;
   if (enableSuspenseServerRenderer) {
     switch (boundaryFiber.tag) {
@@ -2161,6 +2175,21 @@ function warnIfNotCurrentlyActingUpdatesInDEV(fiber: Fiber): void {
   }
 }
 
+export function restorePendingSchedulers(
+  root: FiberRoot,
+  expirationTime: ExpirationTime,
+): void {
+  const pendingSchedulersMap = root.pendingSchedulersMap;
+  let schedulers = pendingSchedulersMap.get(expirationTime);
+  if (schedulers == null) {
+    schedulers = new Set();
+    pendingSchedulersMap.set(expirationTime, schedulers);
+  }
+  root.memoizedSchedulers.forEach(schedulingFiber => {
+    ((schedulers: any): Set<Fiber>).add(schedulingFiber);
+  });
+}
+
 export const warnIfNotCurrentlyActingUpdatesInDev = warnIfNotCurrentlyActingUpdatesInDEV;
 
 let componentsWithSuspendedDiscreteUpdates = null;
@@ -2294,22 +2323,23 @@ function startWorkOnPendingInteraction(root, expirationTime) {
       }
     },
   );
-  const schedulers: Set<Fiber> = new Set();
-  root.pendingSchedulerMap.forEach(
-    (schedulingFibers, scheduledExpirationTime) => {
-      if (scheduledExpirationTime >= expirationTime) {
-        schedulingFibers.forEach(fiber => schedulers.add(fiber));
-      }
-    },
-  );
+
+  const memoizedSchedulers: Set<Fiber> = new Set();
+  const pendingSchedulersMap = root.pendingSchedulersMap;
+  pendingSchedulersMap.forEach((schedulers, scheduledExpirationTime) => {
+    if (scheduledExpirationTime >= expirationTime) {
+      pendingSchedulersMap.delete(scheduledExpirationTime);
+      schedulers.forEach(fiber => memoizedSchedulers.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;
-  root.memoizedSchedulers = schedulers;
+  root.memoizedSchedulers = memoizedSchedulers;
 
   if (interactions.size > 0) {
     const subscriber = __subscriberRef.current;
@@ -2377,12 +2407,5 @@ function finishPendingInteractions(root, committedExpirationTime) {
         }
       },
     );
-
-    const pendingSchedulerMap = root.pendingSchedulerMap;
-    pendingSchedulerMap.forEach((schedulingFibers, scheduledExpirationTime) => {
-      if (scheduledExpirationTime > earliestRemainingTimeAfterCommit) {
-        pendingSchedulerMap.delete(scheduledExpirationTime);
-      }
-    });
   }
 }
diff --git a/packages/react-reconciler/src/ReactFiberUnwindWork.js b/packages/react-reconciler/src/ReactFiberUnwindWork.js
@@ -76,6 +76,7 @@ import {
   pingSuspendedRoot,
   resolveRetryThenable,
   checkForWrongSuspensePriorityInDEV,
+  restorePendingSchedulers,
 } from './ReactFiberScheduler';
 
 import invariant from 'shared/invariant';
@@ -186,6 +187,9 @@ function attachPingListener(
     );
     if (enableSchedulerTracing) {
       ping = Schedule_tracing_wrap(ping);
+
+      // If we have pending work still, restore the original schedulers
+      restorePendingSchedulers(root, renderExpirationTime);
     }
     thenable.then(ping, ping);
   }
@@ -318,6 +322,9 @@ function throwException(
           let retry = resolveRetryThenable.bind(null, workInProgress, thenable);
           if (enableSchedulerTracing) {
             retry = Schedule_tracing_wrap(retry);
+
+            // If we have pending work still, restore the original schedulers
+            restorePendingSchedulers(root, renderExpirationTime);
           }
           thenable.then(retry, retry);
         }