Restart on suspend if return path has an update

Similar to the previous commit, if we suspend with a delay, and
something in the return path has a pending update, we should abort
the current render and switch to the update instead.

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -2064,7 +2064,7 @@ function updateDehydratedSuspenseComponent(
     // render something, if we time out. Even if that requires us to delete everything and
     // skip hydration.
     // Delay having to do this as long as the suspense timeout allows us.
-    renderDidSuspendDelayIfPossible();
+    renderDidSuspendDelayIfPossible(workInProgress);
     return retrySuspenseComponentWithoutHydrating(
       current,
       workInProgress,

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -922,7 +922,7 @@ function completeWork(
           } else {
             // Otherwise, we're going to have to hide content so we should
             // suspend for longer if possible.
-            renderDidSuspendDelayIfPossible();
+            renderDidSuspendDelayIfPossible(workInProgress);
           }
         }
       }

packages/react-reconciler/src/ReactFiberHooks.js

@@ -531,6 +531,14 @@ export function resetHooks(): void {
   // This is used to reset the state of this module when a component throws.
   // It's also called inside mountIndeterminateComponent if we determine the
   // component is a module-style component.
+
+  if (currentlyRenderingFiber !== null) {
+    // Even though this component didn't complete, set the remaining time left
+    // on this fiber. This is sometimes useful when suspending to determine if
+    // there's a lower priority update that could "unsuspend."
+    currentlyRenderingFiber.expirationTime = remainingExpirationTime;
+  }
+
   renderExpirationTime = NoWork;
   currentlyRenderingFiber = null;
 

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -956,6 +956,8 @@ function renderRoot(
           thrownValue,
           renderExpirationTime,
         );
+        // TODO: This is not wrapped in a try-catch, so if the complete phase
+        // throws, we won't capture it.
         workInProgress = completeUnitOfWork(sourceFiber);
       }
     } while (true);
@@ -1259,13 +1261,37 @@ export function renderDidSuspend(): void {
   }
 }
 
-export function renderDidSuspendDelayIfPossible(): void {
+export function renderDidSuspendDelayIfPossible(suspendedWork: Fiber): void {
   if (
     workInProgressRootExitStatus === RootIncomplete ||
     workInProgressRootExitStatus === RootSuspended
   ) {
     workInProgressRootExitStatus = RootSuspendedWithDelay;
   }
+
+  if (workInProgressRoot !== null) {
+    // Check if the component that suspsended, or any components in the return
+    // path, have a pending update. If so, those updates might unsuspend us, so
+    // interrupt the current render and restart.
+    let nextAfterSuspendedTime = NoWork;
+    let fiber = suspendedWork;
+    while (fiber !== null) {
+      const updateExpirationTime = fiber.expirationTime;
+      if (updateExpirationTime > nextAfterSuspendedTime) {
+        nextAfterSuspendedTime = updateExpirationTime;
+      }
+      fiber = fiber.return;
+    }
+
+    if (nextAfterSuspendedTime !== NoWork) {
+      // Mark the current render as suspended, and then mark that there's a
+      // pending update.
+      // TODO: This should immediately interrupt the current render, instead
+      // of waiting until the next time we yield.
+      markRootSuspendedAtTime(workInProgressRoot, renderExpirationTime);
+      markRootUpdatedAtTime(workInProgressRoot, nextAfterSuspendedTime);
+    }
+  }
 }
 
 export function renderDidError() {

packages/react-reconciler/src/__tests__/ReactSuspense-test.internal.js

@@ -315,6 +315,127 @@ describe('ReactSuspense', () => {
     },
   );
 
+  it(
+    'interrupts current render when something suspends with a ' +
+      "delay and we've already skipped over a lower priority update in " +
+      'a parent',
+    () => {
+      function interrupt() {
+        // React has a heuristic to batch all updates that occur within the same
+        // event. This is a trick to circumvent that heuristic.
+        ReactTestRenderer.create('whatever');
+      }
+
+      function App({shouldSuspend, step}) {
+        return (
+          <>
+            <Text text={`A${step}`} />
+            <Suspense fallback={<Text text="Loading..." />}>
+              {shouldSuspend ? <AsyncText text="Async" ms={2000} /> : null}
+            </Suspense>
+            <Text text={`B${step}`} />
+            <Text text={`C${step}`} />
+          </>
+        );
+      }
+
+      const root = ReactTestRenderer.create(null, {
+        unstable_isConcurrent: true,
+      });
+
+      root.update(<App shouldSuspend={false} step={0} />);
+      expect(Scheduler).toFlushAndYield(['A0', 'B0', 'C0']);
+      expect(root).toMatchRenderedOutput('A0B0C0');
+
+      // This update will suspend.
+      root.update(<App shouldSuspend={true} step={1} />);
+
+      // Need to move into the next async bucket.
+      Scheduler.unstable_advanceTime(1000);
+      // Do a bit of work, then interrupt to trigger a restart.
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+      interrupt();
+
+      // Schedule another update. This will have lower priority because of
+      // the interrupt trick above.
+      root.update(<App shouldSuspend={false} step={2} />);
+
+      expect(Scheduler).toFlushAndYieldThrough([
+        // Should have restarted the first update, because of the interruption
+        'A1',
+        'Suspend! [Async]',
+        'Loading...',
+        'B1',
+      ]);
+
+      // Should not have committed loading state
+      expect(root).toMatchRenderedOutput('A0B0C0');
+
+      // After suspending, should abort the first update and switch to the
+      // second update. So, C1 should not appear in the log.
+      // TODO: This should work even if React does not yield to the main
+      // thread. Should use same mechanism as selective hydration to interrupt
+      // the render before the end of the current slice of work.
+      expect(Scheduler).toFlushAndYield(['A2', 'B2', 'C2']);
+
+      expect(root).toMatchRenderedOutput('A2B2C2');
+    },
+  );
+
+  it(
+    'interrupts current render when something suspends with a ' +
+      'delay, and a parent received an update after it completed',
+    () => {
+      function App({shouldSuspend, step}) {
+        return (
+          <>
+            <Text text={`A${step}`} />
+            <Suspense fallback={<Text text="Loading..." />}>
+              {shouldSuspend ? <AsyncText text="Async" ms={2000} /> : null}
+            </Suspense>
+            <Text text={`B${step}`} />
+            <Text text={`C${step}`} />
+          </>
+        );
+      }
+
+      const root = ReactTestRenderer.create(null, {
+        unstable_isConcurrent: true,
+      });
+
+      root.update(<App shouldSuspend={false} step={0} />);
+      expect(Scheduler).toFlushAndYield(['A0', 'B0', 'C0']);
+      expect(root).toMatchRenderedOutput('A0B0C0');
+
+      // This update will suspend.
+      root.update(<App shouldSuspend={true} step={1} />);
+      // Flush past the root, but stop before the async component.
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      // Schedule an update on the root, which already completed.
+      root.update(<App shouldSuspend={false} step={2} />);
+      // We'll keep working on the existing update.
+      expect(Scheduler).toFlushAndYieldThrough([
+        // Now the async component suspends
+        'Suspend! [Async]',
+        'Loading...',
+        'B1',
+      ]);
+
+      // Should not have committed loading state
+      expect(root).toMatchRenderedOutput('A0B0C0');
+
+      // After suspending, should abort the first update and switch to the
+      // second update. So, C1 should not appear in the log.
+      // TODO: This should work even if React does not yield to the main
+      // thread. Should use same mechanism as selective hydration to interrupt
+      // the render before the end of the current slice of work.
+      expect(Scheduler).toFlushAndYield(['A2', 'B2', 'C2']);
+
+      expect(root).toMatchRenderedOutput('A2B2C2');
+    },
+  );
+
   it('mounts a lazy class component in non-concurrent mode', async () => {
     class Class extends React.Component {
       componentDidMount() {