Track the next unprocessed level globally

Instead of backtracking the return path. The main advantage over the
backtracking approach is that we don't have to backtrack from the source
fiber. (The main disadvantages are that it requires another module-level
variable, and that it could include updates from unrelated
sibling paths.)

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -177,6 +177,7 @@ import {
   retryDehydratedSuspenseBoundary,
   scheduleWork,
   renderDidSuspendDelayIfPossible,
+  markUnprocessedUpdateTime,
 } from './ReactFiberWorkLoop';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
@@ -2064,7 +2065,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(workInProgress);
+    renderDidSuspendDelayIfPossible();
     return retrySuspenseComponentWithoutHydrating(
       current,
       workInProgress,
@@ -2708,6 +2709,11 @@ function bailoutOnAlreadyFinishedWork(
     stopProfilerTimerIfRunning(workInProgress);
   }
 
+  const updateExpirationTime = workInProgress.expirationTime;
+  if (updateExpirationTime !== NoWork) {
+    markUnprocessedUpdateTime(updateExpirationTime);
+  }
+
   // Check if the children have any pending work.
   const childExpirationTime = workInProgress.childExpirationTime;
   if (childExpirationTime < renderExpirationTime) {

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(workInProgress);
+            renderDidSuspendDelayIfPossible();
           }
         }
       }

packages/react-reconciler/src/ReactFiberHooks.js

@@ -43,6 +43,7 @@ import {
   warnIfNotCurrentlyActingUpdatesInDev,
   warnIfNotScopedWithMatchingAct,
   markRenderEventTimeAndConfig,
+  markUnprocessedUpdateTime,
 } from './ReactFiberWorkLoop';
 
 import invariant from 'shared/invariant';
@@ -532,13 +533,6 @@ export function resetHooks(): void {
   // 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;
 
@@ -763,6 +757,7 @@ function updateReducer<S, I, A>(
         // Update the remaining priority in the queue.
         if (updateExpirationTime > remainingExpirationTime) {
           remainingExpirationTime = updateExpirationTime;
+          markUnprocessedUpdateTime(remainingExpirationTime);
         }
       } else {
         // This update does have sufficient priority.

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -226,6 +226,10 @@ let workInProgressRootExitStatus: RootExitStatus = RootIncomplete;
 let workInProgressRootLatestProcessedExpirationTime: ExpirationTime = Sync;
 let workInProgressRootLatestSuspenseTimeout: ExpirationTime = Sync;
 let workInProgressRootCanSuspendUsingConfig: null | SuspenseConfig = null;
+// The work left over by components that were visited during this render. Only
+// includes unprocessed updates, not work in bailed out children.
+let workInProgressRootNextUnprocessedUpdateTime: ExpirationTime = NoWork;
+
 // If we're pinged while rendering we don't always restart immediately.
 // This flag determines if it might be worthwhile to restart if an opportunity
 // happens latere.
@@ -484,21 +488,27 @@ function markUpdateTimeFromFiberToRoot(fiber, expirationTime) {
   }
 
   if (root !== null) {
-    if (workInProgressRootExitStatus === RootSuspendedWithDelay) {
-      // The root already suspended with a delay, which means this render
-      // definitely won't finish. Since we have a new update, let's mark it as
-      // suspended now, right before marking the incoming update. This has the
-      // effect of interrupting the current render and switching to the update.
-      // TODO: This happens to work when receiving an update during the render
-      // phase, because of the trick inside computeExpirationForFiber to
-      // subtract 1 from `renderExpirationTime` to move it into a
-      // separate bucket. But we should probably model it with an exception,
-      // using the same mechanism we use to force hydration of a subtree.
-      // TODO: This does not account for low pri updates that were already
-      // scheduled before the root started rendering. Need to track the next
-      // pending expiration time (perhaps by backtracking the return path) and
-      // then trigger a restart in the `renderDidSuspendDelayIfPossible` path.
-      markRootSuspendedAtTime(root, renderExpirationTime);
+    if (workInProgressRoot === root) {
+      // Received an update to a tree that's in the middle of rendering. Mark
+      // that's unprocessed work on this root.
+      markUnprocessedUpdateTime(expirationTime);
+
+      if (workInProgressRootExitStatus === RootSuspendedWithDelay) {
+        // The root already suspended with a delay, which means this render
+        // definitely won't finish. Since we have a new update, let's mark it as
+        // suspended now, right before marking the incoming update. This has the
+        // effect of interrupting the current render and switching to the update.
+        // TODO: This happens to work when receiving an update during the render
+        // phase, because of the trick inside computeExpirationForFiber to
+        // subtract 1 from `renderExpirationTime` to move it into a
+        // separate bucket. But we should probably model it with an exception,
+        // using the same mechanism we use to force hydration of a subtree.
+        // TODO: This does not account for low pri updates that were already
+        // scheduled before the root started rendering. Need to track the next
+        // pending expiration time (perhaps by backtracking the return path) and
+        // then trigger a restart in the `renderDidSuspendDelayIfPossible` path.
+        markRootSuspendedAtTime(root, renderExpirationTime);
+      }
     }
     // Mark that the root has a pending update.
     markRootUpdatedAtTime(root, expirationTime);
@@ -856,6 +866,7 @@ function prepareFreshStack(root, expirationTime) {
   workInProgressRootLatestProcessedExpirationTime = Sync;
   workInProgressRootLatestSuspenseTimeout = Sync;
   workInProgressRootCanSuspendUsingConfig = null;
+  workInProgressRootNextUnprocessedUpdateTime = NoWork;
   workInProgressRootHasPendingPing = false;
 
   if (enableSchedulerTracing) {
@@ -1255,42 +1266,42 @@ export function markRenderEventTimeAndConfig(
   }
 }
 
+export function markUnprocessedUpdateTime(
+  expirationTime: ExpirationTime,
+): void {
+  if (expirationTime > workInProgressRootNextUnprocessedUpdateTime) {
+    workInProgressRootNextUnprocessedUpdateTime = expirationTime;
+  }
+}
+
 export function renderDidSuspend(): void {
   if (workInProgressRootExitStatus === RootIncomplete) {
     workInProgressRootExitStatus = RootSuspended;
   }
 }
 
-export function renderDidSuspendDelayIfPossible(suspendedWork: Fiber): void {
+export function renderDidSuspendDelayIfPossible(): 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);
-    }
+  // Check if there's a lower priority update somewhere else in the tree.
+  if (
+    workInProgressRootNextUnprocessedUpdateTime !== NoWork &&
+    workInProgressRoot !== null
+  ) {
+    // 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,
+      workInProgressRootNextUnprocessedUpdateTime,
+    );
   }
 }
 

packages/react-reconciler/src/ReactUpdateQueue.js

@@ -103,7 +103,10 @@ import {
 } from 'shared/ReactFeatureFlags';
 
 import {StrictMode} from './ReactTypeOfMode';
-import {markRenderEventTimeAndConfig} from './ReactFiberWorkLoop';
+import {
+  markRenderEventTimeAndConfig,
+  markUnprocessedUpdateTime,
+} from './ReactFiberWorkLoop';
 
 import invariant from 'shared/invariant';
 import warningWithoutStack from 'shared/warningWithoutStack';
@@ -580,6 +583,7 @@ export function processUpdateQueue<State>(
   // dealt with the props. Context in components that specify
   // shouldComponentUpdate is tricky; but we'll have to account for
   // that regardless.
+  markUnprocessedUpdateTime(newExpirationTime);
   workInProgress.expirationTime = newExpirationTime;
   workInProgress.memoizedState = resultState;
 

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

@@ -382,6 +382,101 @@ describe('ReactSuspense', () => {
     },
   );
 
+  it(
+    'interrupts current render when something suspends with a ' +
+      "delay and we've already bailed out lower priority update in " +
+      'a parent',
+    async () => {
+      // This is similar to the previous test case, except this covers when
+      // React completely bails out on the parent component, without processing
+      // the update queue.
+
+      const {useState} = React;
+
+      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');
+      }
+
+      let setShouldSuspend;
+      function Async() {
+        const [shouldSuspend, _setShouldSuspend] = useState(false);
+        setShouldSuspend = _setShouldSuspend;
+        return (
+          <>
+            <Text text="A" />
+            <Suspense fallback={<Text text="Loading..." />}>
+              {shouldSuspend ? <AsyncText text="Async" ms={2000} /> : null}
+            </Suspense>
+            <Text text="B" />
+            <Text text="C" />
+          </>
+        );
+      }
+
+      let setShouldHideInParent;
+      function App() {
+        const [shouldHideInParent, _setShouldHideInParent] = useState(false);
+        setShouldHideInParent = _setShouldHideInParent;
+        Scheduler.unstable_yieldValue(
+          'shouldHideInParent: ' + shouldHideInParent,
+        );
+        return shouldHideInParent ? <Text text="(empty)" /> : <Async />;
+      }
+
+      const root = ReactTestRenderer.create(null, {
+        unstable_isConcurrent: true,
+      });
+
+      await ReactTestRenderer.act(async () => {
+        root.update(<App />);
+        expect(Scheduler).toFlushAndYield([
+          'shouldHideInParent: false',
+          'A',
+          'B',
+          'C',
+        ]);
+        expect(root).toMatchRenderedOutput('ABC');
+
+        // This update will suspend.
+        setShouldSuspend(true);
+
+        // 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(['A']);
+        interrupt();
+        // Should not have committed loading state
+        expect(root).toMatchRenderedOutput('ABC');
+
+        // Schedule another update. This will have lower priority because of
+        // the interrupt trick above.
+        setShouldHideInParent(true);
+
+        expect(Scheduler).toFlushAndYieldThrough([
+          // Should have restarted the first update, because of the interruption
+          'A',
+          'Suspend! [Async]',
+          'Loading...',
+          'B',
+        ]);
+
+        // Should not have committed loading state
+        expect(root).toMatchRenderedOutput('ABC');
+
+        // After suspending, should abort the first update and switch to the
+        // second update.
+        expect(Scheduler).toFlushAndYield([
+          'shouldHideInParent: true',
+          '(empty)',
+        ]);
+
+        expect(root).toMatchRenderedOutput('(empty)');
+      });
+    },
+  );
+
   it(
     'interrupts current render when something suspends with a ' +
       'delay, and a parent received an update after it completed',