Track "pending" and "suspended" ranges

A FiberRoot can have pending work at many distinct priorities. (Note: we
refer to these levels as "expiration times" to distinguish the concept
from Scheduler's notion of priority levels, which represent broad
categories of work. React expiration times are more granualar. They're
more like a concurrent thread ID, which also happens to correspond to a
moment on a timeline. It's an overloaded concept and I'm handwaving over
some of the details.)

Given a root, there's no convenient way to read all the pending levels
in the entire tree, i.e. there's no single queue-like structure that
tracks all the levels, because that granularity of information is not
needed by our algorithms. Instead we track the subset of information
that we actually need — most importantly, the highest priority level
that exists in the entire tree.

Aside from that, the other information we track includes the range of
pending levels that are known to be suspended, and therefore should not
be worked on.

This is a refactor of how that information is tracked, and what each
field represents:

- A *pending* level is work that is unfinished, or not yet committed.
  This includes work that is suspended from committing.
  `firstPendingTime` and `lastPendingTime` represent the range of
  pending work. (Previously, "pending" was the same as "not suspended.")
- A *suspended* level is work that did not complete because data was
  missing. `firstSuspendedTime` and `lastSuspendedTime` represent the
  range of suspended work. It is a subset of the pending range. (These
  fields are new to this commit.)
- `nextAfterSuspendedTime` represents the next known level that comes
  after the suspended range.

This commit doesn't change much in terms of observable behavior. The one
change is that, when a level is suspended, React will continue working
on the next known level instead of jumping straight to the last pending
level. Subsequent commits will use this new structure for a more
substantial refactor for how tasks are scheduled per root.

Author: acdlite

packages/react-reconciler/src/ReactFiberRoot.js

@@ -73,8 +73,15 @@ type BaseFiberRootProperties = {|
   firstPendingTime: ExpirationTime,
   // The latest pending expiration time that exists in the tree
   lastPendingTime: ExpirationTime,
-  // The time at which a suspended component pinged the root to render again
-  pingTime: ExpirationTime,
+  // The earliest suspended expiration time that exists in the tree
+  firstSuspendedTime: ExpirationTime,
+  // The latest suspended expiration time that exists in the tree
+  lastSuspendedTime: ExpirationTime,
+  // The next known expiration time after the suspended range
+  nextAfterSuspendedTime: ExpirationTime,
+  // The latest time at which a suspended component pinged the root to
+  // render again
+  lastPingedTime: ExpirationTime,
 |};
 
 // The following attributes are only used by interaction tracing builds.
@@ -120,7 +127,10 @@ function FiberRootNode(containerInfo, tag, hydrate) {
   this.callbackExpirationTime = NoWork;
   this.firstPendingTime = NoWork;
   this.lastPendingTime = NoWork;
-  this.pingTime = NoWork;
+  this.firstSuspendedTime = NoWork;
+  this.lastSuspendedTime = NoWork;
+  this.nextAfterSuspendedTime = NoWork;
+  this.lastPingedTime = NoWork;
 
   if (enableSchedulerTracing) {
     this.interactionThreadID = unstable_getThreadID();
@@ -151,3 +161,53 @@ export function createFiberRoot(
 
   return root;
 }
+
+export function isRootSuspendedAtTime(
+  root: FiberRoot,
+  expirationTime: ExpirationTime,
+): boolean {
+  const firstSuspendedTime = root.firstSuspendedTime;
+  const lastSuspendedTime = root.lastSuspendedTime;
+  return (
+    firstSuspendedTime !== NoWork &&
+    (firstSuspendedTime >= expirationTime &&
+      lastSuspendedTime <= expirationTime)
+  );
+}
+
+export function markRootSuspendedAtTime(
+  root: FiberRoot,
+  expirationTime: ExpirationTime,
+): void {
+  const firstSuspendedTime = root.firstSuspendedTime;
+  const lastSuspendedTime = root.lastSuspendedTime;
+  if (firstSuspendedTime < expirationTime) {
+    root.firstSuspendedTime = expirationTime;
+  }
+  if (lastSuspendedTime > expirationTime || firstSuspendedTime === NoWork) {
+    root.lastSuspendedTime = expirationTime;
+  }
+
+  if (expirationTime <= root.lastPingedTime) {
+    root.lastPingedTime = NoWork;
+  }
+}
+
+export function markRootUnsuspendedAtTime(
+  root: FiberRoot,
+  expirationTime: ExpirationTime,
+): void {
+  if (expirationTime <= root.lastSuspendedTime) {
+    // The entire suspended range is now unsuspended.
+    root.firstSuspendedTime = root.lastSuspendedTime = root.nextAfterSuspendedTime = NoWork;
+  } else if (expirationTime <= root.firstSuspendedTime) {
+    // Part of the suspended range is now unsuspended. Narrow the range to
+    // include everything between the unsuspended time (non-inclusive) and the
+    // last suspended time.
+    root.firstSuspendedTime = expirationTime - 1;
+  }
+
+  if (expirationTime <= root.lastPingedTime) {
+    root.lastPingedTime = NoWork;
+  }
+}

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -66,6 +66,11 @@ import {
 } from './ReactFiberHostConfig';
 
 import {createWorkInProgress, assignFiberPropertiesInDEV} from './ReactFiber';
+import {
+  isRootSuspendedAtTime,
+  markRootSuspendedAtTime,
+  markRootUnsuspendedAtTime,
+} from './ReactFiberRoot';
 import {
   NoMode,
   StrictMode,
@@ -377,8 +382,6 @@ export function scheduleUpdateOnFiber(
     return;
   }
 
-  root.pingTime = NoWork;
-
   checkForInterruption(fiber, expirationTime);
   recordScheduleUpdate();
 
@@ -492,6 +495,9 @@ function markUpdateTimeFromFiberToRoot(fiber, expirationTime) {
     if (lastPendingTime === NoWork || expirationTime < lastPendingTime) {
       root.lastPendingTime = expirationTime;
     }
+
+    // Mark that the root is no longer suspended at this time.
+    markRootUnsuspendedAtTime(root, expirationTime);
   }
 
   return root;
@@ -807,13 +813,6 @@ function renderRoot(
     'Should not already be working.',
   );
 
-  if (root.firstPendingTime < expirationTime) {
-    // If there's no work left at this expiration time, exit immediately. This
-    // happens when multiple callbacks are scheduled for a single root, but an
-    // earlier callback flushes the work of a later one.
-    return null;
-  }
-
   if (isSync && root.finishedExpirationTime === expirationTime) {
     // There's already a pending commit at this expiration time.
     // TODO: This is poorly factored. This case only exists for the
@@ -831,21 +830,25 @@ function renderRoot(
   } else if (workInProgressRootExitStatus === RootSuspendedWithDelay) {
     // We could've received an update at a lower priority while we yielded.
     // We're suspended in a delayed state. Once we complete this render we're
-    // just going to try to recover at the last pending time anyway so we might
-    // as well start doing that eagerly.
+    // just going to try to recover at the pending time anyway so we might as
+    // well start doing that eagerly.
+    //
     // Ideally we should be able to do this even for retries but we don't yet
     // know if we're going to process an update which wants to commit earlier,
     // and this path happens very early so it would happen too often. Instead,
     // for that case, we'll wait until we complete.
     if (workInProgressRootHasPendingPing) {
       // We have a ping at this expiration. Let's restart to see if we get unblocked.
       prepareFreshStack(root, expirationTime);
-    } else {
-      const lastPendingTime = root.lastPendingTime;
-      if (lastPendingTime < expirationTime) {
-        // There's lower priority work. It might be unsuspended. Try rendering
-        // at that level immediately, while preserving the position in the queue.
-        return renderRoot.bind(null, root, lastPendingTime);
+    } else if (!isSync) {
+      // Check if there's work that isn't in the suspended range
+      const firstPendingTime = root.firstPendingTime;
+      if (!isRootSuspendedAtTime(root, firstPendingTime)) {
+        // There's a pending update that falls outside the range of
+        // suspended work.
+        if (firstPendingTime > expirationTime) {
+          return renderRoot.bind(null, root, firstPendingTime);
+        }
       }
     }
   }
@@ -958,7 +961,8 @@ function renderRoot(
   // something suspended, wait to commit it after a timeout.
   stopFinishedWorkLoopTimer();
 
-  root.finishedWork = root.current.alternate;
+  const finishedWork: Fiber = ((root.finishedWork =
+    root.current.alternate): any);
   root.finishedExpirationTime = expirationTime;
 
   const isLocked = resolveLocksOnRoot(root, expirationTime);
@@ -1002,6 +1006,11 @@ function renderRoot(
       return commitRoot.bind(null, root);
     }
     case RootSuspended: {
+      markRootSuspendedAtTime(root, expirationTime);
+      const lastSuspendedTime = root.lastSuspendedTime;
+      if (expirationTime === lastSuspendedTime) {
+        root.nextAfterSuspendedTime = getRemainingExpirationTime(finishedWork);
+      }
       flushSuspensePriorityWarningInDEV();
 
       // We have an acceptable loading state. We need to figure out if we should
@@ -1038,11 +1047,20 @@ function renderRoot(
             prepareFreshStack(root, expirationTime);
             return renderRoot.bind(null, root, expirationTime);
           }
-          const lastPendingTime = root.lastPendingTime;
-          if (lastPendingTime < expirationTime) {
+
+          const nextAfterSuspendedTime = root.nextAfterSuspendedTime;
+          if (nextAfterSuspendedTime !== NoWork) {
             // There's lower priority work. It might be unsuspended. Try rendering
             // at that level.
-            return renderRoot.bind(null, root, lastPendingTime);
+            return renderRoot.bind(null, root, nextAfterSuspendedTime);
+          }
+          if (
+            lastSuspendedTime !== NoWork &&
+            lastSuspendedTime !== expirationTime
+          ) {
+            // We should prefer to render the fallback of at the last
+            // suspended level.
+            return renderRoot.bind(null, root, lastSuspendedTime);
           }
           // The render is suspended, it hasn't timed out, and there's no lower
           // priority work to do. Instead of committing the fallback
@@ -1058,6 +1076,11 @@ function renderRoot(
       return commitRoot.bind(null, root);
     }
     case RootSuspendedWithDelay: {
+      markRootSuspendedAtTime(root, expirationTime);
+      const lastSuspendedTime = root.lastSuspendedTime;
+      if (expirationTime === lastSuspendedTime) {
+        root.nextAfterSuspendedTime = getRemainingExpirationTime(finishedWork);
+      }
       flushSuspensePriorityWarningInDEV();
 
       if (
@@ -1077,11 +1100,20 @@ function renderRoot(
           prepareFreshStack(root, expirationTime);
           return renderRoot.bind(null, root, expirationTime);
         }
-        const lastPendingTime = root.lastPendingTime;
-        if (lastPendingTime < expirationTime) {
+
+        const nextAfterSuspendedTime = root.nextAfterSuspendedTime;
+        if (nextAfterSuspendedTime !== NoWork) {
           // There's lower priority work. It might be unsuspended. Try rendering
-          // at that level immediately.
-          return renderRoot.bind(null, root, lastPendingTime);
+          // at that level.
+          return renderRoot.bind(null, root, nextAfterSuspendedTime);
+        }
+        if (
+          lastSuspendedTime !== NoWork &&
+          lastSuspendedTime !== expirationTime
+        ) {
+          // We should prefer to render the fallback of at the last
+          // suspended level.
+          return renderRoot.bind(null, root, lastSuspendedTime);
         }
 
         let msUntilTimeout;
@@ -1425,6 +1457,14 @@ function completeUnitOfWork(unitOfWork: Fiber): Fiber | null {
   return null;
 }
 
+function getRemainingExpirationTime(fiber: Fiber) {
+  const updateExpirationTime = fiber.expirationTime;
+  const childExpirationTime = fiber.childExpirationTime;
+  return updateExpirationTime > childExpirationTime
+    ? updateExpirationTime
+    : childExpirationTime;
+}
+
 function resetChildExpirationTime(completedWork: Fiber) {
   if (
     renderExpirationTime !== Never &&
@@ -1540,19 +1580,19 @@ function commitRootImpl(root, renderPriorityLevel) {
 
   // Update the first and last pending times on this root. The new first
   // pending time is whatever is left on the root fiber.
-  const updateExpirationTimeBeforeCommit = finishedWork.expirationTime;
-  const childExpirationTimeBeforeCommit = finishedWork.childExpirationTime;
-  const firstPendingTimeBeforeCommit =
-    childExpirationTimeBeforeCommit > updateExpirationTimeBeforeCommit
-      ? childExpirationTimeBeforeCommit
-      : updateExpirationTimeBeforeCommit;
-  root.firstPendingTime = firstPendingTimeBeforeCommit;
-  if (firstPendingTimeBeforeCommit < root.lastPendingTime) {
+  const remainingExpirationTimeBeforeCommit = getRemainingExpirationTime(
+    finishedWork,
+  );
+  root.firstPendingTime = remainingExpirationTimeBeforeCommit;
+  if (remainingExpirationTimeBeforeCommit < root.lastPendingTime) {
     // This usually means we've finished all the work, but it can also happen
     // when something gets downprioritized during render, like a hidden tree.
-    root.lastPendingTime = firstPendingTimeBeforeCommit;
+    root.lastPendingTime = remainingExpirationTimeBeforeCommit;
   }
 
+  // Mark that the root is no longer suspended at the finished time
+  markRootUnsuspendedAtTime(root, expirationTime);
+
   if (root === workInProgressRoot) {
     // We can reset these now that they are finished.
     workInProgressRoot = null;
@@ -2148,20 +2188,19 @@ export function pingSuspendedRoot(
     return;
   }
 
-  const lastPendingTime = root.lastPendingTime;
-  if (lastPendingTime < suspendedTime) {
+  if (!isRootSuspendedAtTime(root, suspendedTime)) {
     // The root is no longer suspended at this time.
     return;
   }
 
-  const pingTime = root.pingTime;
-  if (pingTime !== NoWork && pingTime < suspendedTime) {
+  const lastPingedTime = root.lastPingedTime;
+  if (lastPingedTime !== NoWork && lastPingedTime < suspendedTime) {
     // There's already a lower priority ping scheduled.
     return;
   }
 
   // Mark the time at which this ping was scheduled.
-  root.pingTime = suspendedTime;
+  root.lastPingedTime = suspendedTime;
 
   if (root.finishedExpirationTime === suspendedTime) {
     // If there's a pending fallback waiting to commit, throw it away.

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

@@ -518,6 +518,71 @@ describe('ReactSuspenseWithNoopRenderer', () => {
     expect(ReactNoop.getChildren()).toEqual([span('(empty)')]);
   });
 
+  it('tries each subsequent level after suspending', async () => {
+    const root = ReactNoop.createRoot();
+
+    function App({step, shouldSuspend}) {
+      return (
+        <Suspense fallback="Loading...">
+          <Text text="Sibling" />
+          {shouldSuspend ? (
+            <AsyncText ms={10000} text={'Step ' + step} />
+          ) : (
+            <Text text={'Step ' + step} />
+          )}
+        </Suspense>
+      );
+    }
+
+    function interrupt() {
+      // React has a heuristic to batch all updates that occur within the same
+      // event. This is a trick to circumvent that heuristic.
+      ReactNoop.flushSync(() => {
+        ReactNoop.renderToRootWithID(null, 'other-root');
+      });
+    }
+
+    // Mount the Suspense boundary without suspending, so that the subsequent
+    // updates suspend with a delay.
+    await ReactNoop.act(async () => {
+      root.render(<App step={0} shouldSuspend={false} />);
+    });
+    await advanceTimers(1000);
+    expect(Scheduler).toHaveYielded(['Sibling', 'Step 0']);
+
+    // Schedule an update at several distinct expiration times
+    await ReactNoop.act(async () => {
+      root.render(<App step={1} shouldSuspend={true} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['Sibling']);
+      interrupt();
+
+      root.render(<App step={2} shouldSuspend={true} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['Sibling']);
+      interrupt();
+
+      root.render(<App step={3} shouldSuspend={true} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['Sibling']);
+      interrupt();
+
+      root.render(<App step={4} shouldSuspend={false} />);
+    });
+
+    // Should suspend at each distinct level
+    expect(Scheduler).toHaveYielded([
+      'Sibling',
+      'Suspend! [Step 1]',
+      'Sibling',
+      'Suspend! [Step 2]',
+      'Sibling',
+      'Suspend! [Step 3]',
+      'Sibling',
+      'Step 4',
+    ]);
+  });
+
   it('forces an expiration after an update times out', async () => {
     ReactNoop.render(
       <Fragment>