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.

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -98,6 +98,8 @@ import {
 } from './ReactUpdateQueue';
 import {
   NoWork,
+  Idle,
+  ContinuousHydration,
   Never,
   Sync,
   computeAsyncExpiration,
@@ -1590,6 +1592,42 @@ 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 a really low
+    // priority that includes all the updates.
+    // TODO: If expirationTime were a bitmask where each bit represents a
+    // separate task thread, this would be: currentChildBits & ~renderBits
+    // TODO: We used to track the lowest pending level for the whole root, but
+    // we removed it to simplify the implementation. It might be worth adding
+    // it back for this use case, to avoid redundant passes.
+    if (renderExpirationTime > ContinuousHydration) {
+      // First we try ContinuousHydration, so that we don't get grouped
+      // with Idle.
+      return ContinuousHydration;
+    } else if (renderExpirationTime > Idle) {
+      // Then we'll try Idle.
+      return Idle;
+    } else {
+      // Already at lowest possible update level.
+      return NoWork;
+    }
+  }
+  // In legacy mode, there's no work left.
+  return NoWork;
+}
+
 function updateSuspenseComponent(
   current,
   workInProgress,
@@ -1831,8 +1869,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 +1940,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 +2039,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;

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

@@ -734,12 +734,22 @@ describe('ReactHooksWithNoopRenderer', () => {
         root.render(<Foo signal={false} />);
         setLabel('B');
       });
-      expect(Scheduler).toHaveYielded(['Suspend!']);
+      expect(Scheduler).toHaveYielded([
+        'Suspend!',
+        // TODO: This second attempt is because React isn't sure if there's
+        // another update at a lower priority. Ideally we wouldn't need it.
+        'Suspend!',
+      ]);
       expect(root).toMatchRenderedOutput(<span prop="A:0" />);
 
       // Rendering again should suspend again.
       root.render(<Foo signal={false} />);
-      expect(Scheduler).toFlushAndYield(['Suspend!']);
+      expect(Scheduler).toFlushAndYield([
+        'Suspend!',
+        // TODO: This second attempt is because React isn't sure if there's
+        // another update at a lower priority. Ideally we wouldn't need it.
+        'Suspend!',
+      ]);
 
       // Flip the signal back to "cancel" the update. However, the update to
       // label should still proceed. It shouldn't have been dropped.

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

@@ -956,7 +956,14 @@ describe('ReactSuspense', () => {
 
       // Update that suspends
       instance.setState({step: 2});
-      expect(Scheduler).toFlushAndYield(['Suspend! [Step: 2]', 'Loading...']);
+      expect(Scheduler).toFlushAndYield([
+        'Suspend! [Step: 2]',
+        'Loading...',
+        // TODO: This second attempt is because React isn't sure if there's
+        // another update at a lower priority. Ideally we wouldn't need it.
+        'Suspend! [Step: 2]',
+        'Loading...',
+      ]);
       jest.advanceTimersByTime(500);
       expect(root).toMatchRenderedOutput('Loading...');
 

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

@@ -2842,6 +2842,7 @@ describe('ReactSuspenseWithNoopRenderer', () => {
       foo.setState({suspend: false});
     });
 
+    expect(Scheduler).toHaveYielded(['Foo']);
     expect(root).toMatchRenderedOutput(<span prop="Foo" />);
   });
 
@@ -2957,4 +2958,142 @@ describe('ReactSuspenseWithNoopRenderer', () => {
       </>,
     );
   });
+
+  it(
+    'multiple updates originating inside a Suspense boundary at different ' +
+      'priority levels are not dropped',
+    async () => {
+      const {useState} = React;
+      const root = ReactNoop.createRoot();
+
+      function Parent() {
+        return (
+          <>
+            <Suspense fallback={<Text text="Loading..." />}>
+              <Child />
+            </Suspense>
+          </>
+        );
+      }
+
+      let setText;
+      let setShouldHide;
+      function Child() {
+        const [text, _setText] = useState('A');
+        const [shouldHide, _setShouldHide] = useState(false);
+        setText = _setText;
+        setShouldHide = _setShouldHide;
+        return shouldHide ? <Text text="(empty)" /> : <AsyncText text={text} />;
+      }
+
+      await ReactNoop.act(async () => {
+        root.render(<Parent />);
+      });
+      expect(Scheduler).toHaveYielded(['Suspend! [A]', 'Loading...']);
+      expect(root).toMatchRenderedOutput(<span prop="Loading..." />);
+
+      await resolveText('A');
+      expect(Scheduler).toHaveYielded(['Promise resolved [A]']);
+      expect(Scheduler).toFlushAndYield(['A']);
+      expect(root).toMatchRenderedOutput(<span prop="A" />);
+
+      await ReactNoop.act(async () => {
+        // Schedule two updates that originate inside the Suspense boundary.
+        // The first one causes the boundary to suspend. The second one is at
+        // lower priority and unsuspends it by hiding the async component.
+        ReactNoop.discreteUpdates(() => {
+          setText('B');
+        });
+        setShouldHide(true);
+      });
+      expect(Scheduler).toHaveYielded([
+        // First we attempt the high pri update. It suspends.
+        'Suspend! [B]',
+        'Loading...',
+        // Then we attempt the low pri update, which finishes successfully.
+        '(empty)',
+      ]);
+      expect(root).toMatchRenderedOutput(<span prop="(empty)" />);
+    },
+  );
+
+  it(
+    'multiple updates originating inside a Suspense boundary at different ' +
+      'priority levels are not dropped, including Idle updates',
+    async () => {
+      const {useState} = React;
+      const root = ReactNoop.createRoot();
+
+      function Parent() {
+        return (
+          <>
+            <Suspense fallback={<Text text="Loading..." />}>
+              <Child />
+            </Suspense>
+          </>
+        );
+      }
+
+      let setText;
+      let setShouldHide;
+      function Child() {
+        const [text, _setText] = useState('A');
+        const [shouldHide, _setShouldHide] = useState(false);
+        setText = _setText;
+        setShouldHide = _setShouldHide;
+        return shouldHide ? <Text text="(empty)" /> : <AsyncText text={text} />;
+      }
+
+      await ReactNoop.act(async () => {
+        root.render(<Parent />);
+      });
+      expect(Scheduler).toHaveYielded(['Suspend! [A]', 'Loading...']);
+      expect(root).toMatchRenderedOutput(<span prop="Loading..." />);
+
+      await resolveText('A');
+      expect(Scheduler).toHaveYielded(['Promise resolved [A]']);
+      expect(Scheduler).toFlushAndYield(['A']);
+      expect(root).toMatchRenderedOutput(<span prop="A" />);
+
+      await ReactNoop.act(async () => {
+        // Schedule two updates that originate inside the Suspense boundary.
+        // The first one causes the boundary to suspend. The second one is at
+        // lower priority and unsuspends it by hiding the async component.
+        setText('B');
+        Scheduler.unstable_runWithPriority(
+          Scheduler.unstable_IdlePriority,
+          () => {
+            setShouldHide(true);
+          },
+        );
+      });
+      expect(Scheduler).toHaveYielded([
+        // First we attempt the high pri update. It suspends.
+        'Suspend! [B]',
+        'Loading...',
+
+        // Then retry at a lower priority level.
+        // NOTE: The current implementation doesn't know which level to attempt,
+        // so it picks ContinuousHydration, which is one level higher than Idle.
+        // Since it doesn't include the Idle update, it will suspend again and
+        // reschedule to try at Idle. If we refactor expiration time to be a
+        // bitmask, we shouldn't need this heuristic.
+        'Suspend! [B]',
+        'Loading...',
+      ]);
+
+      // Commit the placeholder to unblock the Idle update.
+      await advanceTimers(250);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span hidden={true} prop="A" />
+          <span prop="Loading..." />
+        </>,
+      );
+
+      // Now flush the remaining work. The Idle update successfully finishes.
+      expect(Scheduler).toFlushAndYield(['(empty)']);
+      expect(root).toMatchRenderedOutput(<span prop="(empty)" />);
+    },
+  );
 });