[Bugfix] Dropped updates inside a suspended tree (#18384)

* Minor test refactor: `resolveText`

Adds a `resolveText` method as an alternative to using timers. Also
removes dependency on react-cache (for just this one test file; can do
the others later).

Timer option is still there if you provide a `ms` prop.

* Bugfix: Dropped updates in suspended tree

When there are multiple updates at different priority levels inside
a suspended subtree, all but the highest priority one is dropped after
the highest one suspends.

We do have tests that cover this for updates that originate outside of
the Suspense boundary, but not for updates that originate inside.

I'm surprised it's taken us this long to find this issue, but it makes
sense in that transition updates usually originate outside the boundary
or "seam" of the part of the UI that is transitioning.

* Bugfix: Suspense fragment skipped by setState

Fixes a bug where updates inside a suspended tree are dropped because
the fragment fiber we insert to wrap the hidden children is not part of
the return path, so it doesn't get marked during setState.

As a workaround, I recompute `childExpirationTime` right before deciding
to bail out by bubbling it up from the next level of children.

This is something we should consider addressing when we refactor the
Fiber data structure.

* Add back `lastPendingTime` field

This reverts commit 9a54113.

I want to use this so we can check if there might be any lower priority
updates in a suspended tree.

We can remove it again during the expiration times refactor.

* Use `lastPendingTime` instead of Idle

We don't currently have an mechanism to check if there are lower
priority updates in a subtree, but we can check if there are any in the
whole root. This still isn't perfect but it's better than using Idle,
which frequently leads to redundant re-renders.

When we refactor `expirationTime` to be a bitmask, this will no longer
be necessary because we'll know exactly which "task bits" remain.

* Add a test for updating the fallback

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -179,6 +179,7 @@ import {
   scheduleUpdateOnFiber,
   renderDidSuspendDelayIfPossible,
   markUnprocessedUpdateTime,
+  getWorkInProgressRoot,
 } from './ReactFiberWorkLoop';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
@@ -1590,6 +1591,35 @@ function shouldRemainOnFallback(
   );
 }
 
+function getRemainingWorkInPrimaryTree(
+  workInProgress,
+  currentChildExpirationTime,
+  renderExpirationTime,
+) {
+  if (currentChildExpirationTime < renderExpirationTime) {
+    // The highest priority remaining work is not part of this render. So the
+    // remaining work has not changed.
+    return currentChildExpirationTime;
+  }
+  if ((workInProgress.mode & BlockingMode) !== NoMode) {
+    // The highest priority remaining work is part of this render. Since we only
+    // keep track of the highest level, we don't know if there's a lower
+    // priority level scheduled. As a compromise, we'll render at the lowest
+    // known level in the entire tree, since that will include everything.
+    // TODO: If expirationTime were a bitmask where each bit represents a
+    // separate task thread, this would be: currentChildBits & ~renderBits
+    const root = getWorkInProgressRoot();
+    if (root !== null) {
+      const lastPendingTime = root.lastPendingTime;
+      if (lastPendingTime < renderExpirationTime) {
+        return lastPendingTime;
+      }
+    }
+  }
+  // In legacy mode, there's no work left.
+  return NoWork;
+}
+
 function updateSuspenseComponent(
   current,
   workInProgress,
@@ -1831,8 +1861,15 @@ function updateSuspenseComponent(
             fallbackChildFragment.return = workInProgress;
             primaryChildFragment.sibling = fallbackChildFragment;
             fallbackChildFragment.effectTag |= Placement;
-            primaryChildFragment.childExpirationTime = NoWork;
-
+            primaryChildFragment.childExpirationTime = getRemainingWorkInPrimaryTree(
+              workInProgress,
+              // This argument represents the remaining work in the current
+              // primary tree. Since the current tree did not already time out
+              // the direct parent of the primary children is the Suspense
+              // fiber, not a fragment.
+              current.childExpirationTime,
+              renderExpirationTime,
+            );
             workInProgress.memoizedState = SUSPENDED_MARKER;
             workInProgress.child = primaryChildFragment;
 
@@ -1895,6 +1932,11 @@ function updateSuspenseComponent(
         );
         fallbackChildFragment.return = workInProgress;
         primaryChildFragment.sibling = fallbackChildFragment;
+        primaryChildFragment.childExpirationTime = getRemainingWorkInPrimaryTree(
+          workInProgress,
+          currentPrimaryChildFragment.childExpirationTime,
+          renderExpirationTime,
+        );
         primaryChildFragment.childExpirationTime = NoWork;
         // Skip the primary children, and continue working on the
         // fallback children.
@@ -1989,7 +2031,15 @@ function updateSuspenseComponent(
         fallbackChildFragment.return = workInProgress;
         primaryChildFragment.sibling = fallbackChildFragment;
         fallbackChildFragment.effectTag |= Placement;
-        primaryChildFragment.childExpirationTime = NoWork;
+        primaryChildFragment.childExpirationTime = getRemainingWorkInPrimaryTree(
+          workInProgress,
+          // This argument represents the remaining work in the current
+          // primary tree. Since the current tree did not already time out
+          // the direct parent of the primary children is the Suspense
+          // fiber, not a fragment.
+          current.childExpirationTime,
+          renderExpirationTime,
+        );
         // Skip the primary children, and continue working on the
         // fallback children.
         workInProgress.memoizedState = SUSPENDED_MARKER;
@@ -3006,6 +3056,42 @@ function beginWork(
                 renderExpirationTime,
               );
             } else {
+              // The primary child fragment does not have pending work marked
+              // on it...
+
+              // ...usually. There's an unfortunate edge case where the fragment
+              // fiber is not part of the return path of the children, so when
+              // an update happens, the fragment doesn't get marked during
+              // setState. This is something we should consider addressing when
+              // we refactor the Fiber data structure. (There's a test with more
+              // details; to find it, comment out the following block and see
+              // which one fails.)
+              //
+              // As a workaround, we need to recompute the `childExpirationTime`
+              // by bubbling it up from the next level of children. This is
+              // based on similar logic in `resetChildExpirationTime`.
+              let primaryChild = primaryChildFragment.child;
+              while (primaryChild !== null) {
+                const childUpdateExpirationTime = primaryChild.expirationTime;
+                const childChildExpirationTime =
+                  primaryChild.childExpirationTime;
+                if (
+                  (childUpdateExpirationTime !== NoWork &&
+                    childUpdateExpirationTime >= renderExpirationTime) ||
+                  (childChildExpirationTime !== NoWork &&
+                    childChildExpirationTime >= renderExpirationTime)
+                ) {
+                  // Found a child with an update with sufficient priority.
+                  // Use the normal path to render the primary children again.
+                  return updateSuspenseComponent(
+                    current,
+                    workInProgress,
+                    renderExpirationTime,
+                  );
+                }
+                primaryChild = primaryChild.sibling;
+              }
+
               pushSuspenseContext(
                 workInProgress,
                 setDefaultShallowSuspenseContext(suspenseStackCursor.current),

packages/react-reconciler/src/ReactFiberRoot.js

@@ -65,6 +65,8 @@ type BaseFiberRootProperties = {|
   callbackPriority: ReactPriorityLevel,
   // The earliest pending expiration time that exists in the tree
   firstPendingTime: ExpirationTime,
+  // The latest pending expiration time that exists in the tree
+  lastPendingTime: ExpirationTime,
   // The earliest suspended expiration time that exists in the tree
   firstSuspendedTime: ExpirationTime,
   // The latest suspended expiration time that exists in the tree
@@ -122,6 +124,7 @@ function FiberRootNode(containerInfo, tag, hydrate) {
   this.callbackNode = null;
   this.callbackPriority = NoPriority;
   this.firstPendingTime = NoWork;
+  this.lastPendingTime = NoWork;
   this.firstSuspendedTime = NoWork;
   this.lastSuspendedTime = NoWork;
   this.nextKnownPendingLevel = NoWork;
@@ -205,6 +208,10 @@ export function markRootUpdatedAtTime(
   if (expirationTime > firstPendingTime) {
     root.firstPendingTime = expirationTime;
   }
+  const lastPendingTime = root.lastPendingTime;
+  if (lastPendingTime === NoWork || expirationTime < lastPendingTime) {
+    root.lastPendingTime = expirationTime;
+  }
 
   // Update the range of suspended times. Treat everything lower priority or
   // equal to this update as unsuspended.
@@ -232,6 +239,11 @@ export function markRootFinishedAtTime(
 ): void {
   // Update the range of pending times
   root.firstPendingTime = remainingExpirationTime;
+  if (remainingExpirationTime < 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 = remainingExpirationTime;
+  }
 
   // Update the range of suspended times. Treat everything higher priority or
   // equal to this update as unsuspended.

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

@@ -293,7 +293,13 @@ describe('ReactSuspenseList', () => {
 
     await C.resolve();
 
-    expect(Scheduler).toFlushAndYield(['C']);
+    expect(Scheduler).toFlushAndYield([
+      // TODO: Ideally we wouldn't have to retry B. This is an implementation
+      // trade off.
+      'Suspend! [B]',
+
+      'C',
+    ]);
 
     expect(ReactNoop).toMatchRenderedOutput(
       <>