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.

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -3064,6 +3064,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/__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(
       <>

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

@@ -1335,6 +1335,7 @@ describe('ReactSuspenseWithNoopRenderer', () => {
         'Suspend! [Hi]',
         'Loading...',
         // Re-render due to lifecycle update
+        'Suspend! [Hi]',
         'Loading...',
       ]);
       expect(ReactNoop.getChildren()).toEqual([span('Loading...')]);
@@ -3115,4 +3116,89 @@ describe('ReactSuspenseWithNoopRenderer', () => {
       expect(root).toMatchRenderedOutput(<span prop="C" />);
     },
   );
+
+  it(
+    'regression: primary fragment fiber is not always part of setState ' +
+      'return path',
+    async () => {
+      // Reproduces 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.
+      const {useState} = React;
+      const root = ReactNoop.createRoot();
+
+      function Parent() {
+        return (
+          <>
+            <Suspense fallback={<Text text="Loading..." />}>
+              <Child />
+            </Suspense>
+          </>
+        );
+      }
+
+      let setText;
+      function Child() {
+        const [text, _setText] = useState('A');
+        setText = _setText;
+        return <AsyncText text={text} />;
+      }
+
+      // Mount an initial tree. Resolve A so that it doesn't suspend.
+      await resolveText('A');
+      await ReactNoop.act(async () => {
+        root.render(<Parent />);
+      });
+      expect(Scheduler).toHaveYielded(['A']);
+      // At this point, the setState return path follows current fiber.
+      expect(root).toMatchRenderedOutput(<span prop="A" />);
+
+      // Schedule another update. This will "flip" the alternate pairs.
+      await resolveText('B');
+      await ReactNoop.act(async () => {
+        setText('B');
+      });
+      expect(Scheduler).toHaveYielded(['B']);
+      // Now the setState return path follows the *alternate* fiber.
+      expect(root).toMatchRenderedOutput(<span prop="B" />);
+
+      // Schedule another update. This time, we'll suspend.
+      await ReactNoop.act(async () => {
+        setText('C');
+      });
+      expect(Scheduler).toHaveYielded([
+        'Suspend! [C]',
+        'Loading...',
+
+        'Suspend! [C]',
+        'Loading...',
+      ]);
+
+      // Commit. This will insert a fragment fiber to wrap around the component
+      // that triggered the update.
+      await ReactNoop.act(async () => {
+        await advanceTimers(250);
+      });
+      expect(Scheduler).toHaveYielded(['Suspend! [C]']);
+      // The fragment fiber is part of the current tree, but the setState return
+      // path still follows the alternate path. That means the fragment fiber is
+      // not part of the return path.
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span hidden={true} prop="B" />
+          <span prop="Loading..." />
+        </>,
+      );
+
+      // Update again. This should unsuspend the tree.
+      await resolveText('D');
+      await ReactNoop.act(async () => {
+        setText('D');
+      });
+      // Even though the fragment fiber is not part of the return path, we should
+      // be able to finish rendering.
+      expect(Scheduler).toHaveYielded(['D']);
+      expect(root).toMatchRenderedOutput(<span prop="D" />);
+    },
+  );
 });