Bugfix: Do not unhide a suspended tree without finishing the suspended update (#18411)

* Bugfix: Suspended update must finish to unhide

When we commit a fallback, we cannot unhide the content without including
the level that originally suspended. That's because the work at level
outside the boundary (i.e. everything that wasn't hidden during that
render) already committed.

* Test unblocking with a high-pri update

packages/react-debug-tools/src/__tests__/ReactDevToolsHooksIntegration-test.js

@@ -265,6 +265,7 @@ describe('React hooks DevTools integration', () => {
       // Release the lock
       setSuspenseHandler(() => false);
       scheduleUpdate(fiber); // Re-render
+      Scheduler.unstable_flushAll();
       expect(renderer.toJSON().children).toEqual(['Done']);
       scheduleUpdate(fiber); // Re-render
       expect(renderer.toJSON().children).toEqual(['Done']);

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -1558,37 +1558,102 @@ function validateFunctionComponentInDev(workInProgress: Fiber, Component: any) {
   }
 }
 
-const SUSPENDED_MARKER: SuspenseState = {
-  dehydrated: null,
-  retryTime: NoWork,
-};
+function mountSuspenseState(
+  renderExpirationTime: ExpirationTime,
+): SuspenseState {
+  return {
+    dehydrated: null,
+    baseTime: renderExpirationTime,
+    retryTime: NoWork,
+  };
+}
+
+function updateSuspenseState(
+  prevSuspenseState: SuspenseState,
+  renderExpirationTime: ExpirationTime,
+): SuspenseState {
+  const prevSuspendedTime = prevSuspenseState.baseTime;
+  return {
+    dehydrated: null,
+    baseTime:
+      // Choose whichever time is inclusive of the other one. This represents
+      // the union of all the levels that suspended.
+      prevSuspendedTime !== NoWork && prevSuspendedTime < renderExpirationTime
+        ? prevSuspendedTime
+        : renderExpirationTime,
+    retryTime: NoWork,
+  };
+}
 
 function shouldRemainOnFallback(
   suspenseContext: SuspenseContext,
   current: null | Fiber,
   workInProgress: Fiber,
+  renderExpirationTime: ExpirationTime,
 ) {
-  // If the context is telling us that we should show a fallback, and we're not
-  // already showing content, then we should show the fallback instead.
-  return (
-    hasSuspenseContext(
-      suspenseContext,
-      (ForceSuspenseFallback: SuspenseContext),
-    ) &&
-    (current === null || current.memoizedState !== null)
+  // If we're already showing a fallback, there are cases where we need to
+  // remain on that fallback regardless of whether the content has resolved.
+  // For example, SuspenseList coordinates when nested content appears.
+  if (current !== null) {
+    const suspenseState: SuspenseState = current.memoizedState;
+    if (suspenseState !== null) {
+      // Currently showing a fallback. If the current render includes
+      // the level that triggered the fallback, we must continue showing it,
+      // regardless of what the Suspense context says.
+      const baseTime = suspenseState.baseTime;
+      if (baseTime !== NoWork && baseTime < renderExpirationTime) {
+        return true;
+      }
+      // Otherwise, fall through to check the Suspense context.
+    } else {
+      // Currently showing content. Don't hide it, even if ForceSuspenseFallack
+      // is true. More precise name might be "ForceRemainSuspenseFallback".
+      // Note: This is a factoring smell. Can't remain on a fallback if there's
+      // no fallback to remain on.
+      return false;
+    }
+  }
+  // Not currently showing content. Consult the Suspense context.
+  return hasSuspenseContext(
+    suspenseContext,
+    (ForceSuspenseFallback: SuspenseContext),
   );
 }
 
 function getRemainingWorkInPrimaryTree(
-  workInProgress,
-  currentChildExpirationTime,
+  current: Fiber,
+  workInProgress: Fiber,
+  currentPrimaryChildFragment: Fiber | null,
   renderExpirationTime,
 ) {
+  const currentParentOfPrimaryChildren =
+    currentPrimaryChildFragment !== null
+      ? currentPrimaryChildFragment
+      : current;
+  const currentChildExpirationTime =
+    currentParentOfPrimaryChildren.childExpirationTime;
+
+  const currentSuspenseState: SuspenseState = current.memoizedState;
+  if (currentSuspenseState !== null) {
+    // This boundary already timed out. Check if this render includes the level
+    // that previously suspended.
+    const baseTime = currentSuspenseState.baseTime;
+    if (
+      baseTime !== NoWork &&
+      baseTime < renderExpirationTime &&
+      baseTime > currentChildExpirationTime
+    ) {
+      // There's pending work at a lower level that might now be unblocked.
+      return baseTime;
+    }
+  }
+
   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
@@ -1630,7 +1695,12 @@ function updateSuspenseComponent(
 
   if (
     didSuspend ||
-    shouldRemainOnFallback(suspenseContext, current, workInProgress)
+    shouldRemainOnFallback(
+      suspenseContext,
+      current,
+      workInProgress,
+      renderExpirationTime,
+    )
   ) {
     // Something in this boundary's subtree already suspended. Switch to
     // rendering the fallback children.
@@ -1746,7 +1816,7 @@ function updateSuspenseComponent(
       primaryChildFragment.sibling = fallbackChildFragment;
       // Skip the primary children, and continue working on the
       // fallback children.
-      workInProgress.memoizedState = SUSPENDED_MARKER;
+      workInProgress.memoizedState = mountSuspenseState(renderExpirationTime);
       workInProgress.child = primaryChildFragment;
       return fallbackChildFragment;
     } else {
@@ -1850,15 +1920,15 @@ function updateSuspenseComponent(
             primaryChildFragment.sibling = fallbackChildFragment;
             fallbackChildFragment.effectTag |= Placement;
             primaryChildFragment.childExpirationTime = getRemainingWorkInPrimaryTree(
+              current,
               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,
+              null,
+              renderExpirationTime,
+            );
+            workInProgress.memoizedState = updateSuspenseState(
+              current.memoizedState,
               renderExpirationTime,
             );
-            workInProgress.memoizedState = SUSPENDED_MARKER;
             workInProgress.child = primaryChildFragment;
 
             // Skip the primary children, and continue working on the
@@ -1921,13 +1991,17 @@ function updateSuspenseComponent(
         fallbackChildFragment.return = workInProgress;
         primaryChildFragment.sibling = fallbackChildFragment;
         primaryChildFragment.childExpirationTime = getRemainingWorkInPrimaryTree(
+          current,
           workInProgress,
-          currentPrimaryChildFragment.childExpirationTime,
+          currentPrimaryChildFragment,
           renderExpirationTime,
         );
         // Skip the primary children, and continue working on the
         // fallback children.
-        workInProgress.memoizedState = SUSPENDED_MARKER;
+        workInProgress.memoizedState = updateSuspenseState(
+          current.memoizedState,
+          renderExpirationTime,
+        );
         workInProgress.child = primaryChildFragment;
         return fallbackChildFragment;
       } else {
@@ -2019,17 +2093,14 @@ function updateSuspenseComponent(
         primaryChildFragment.sibling = fallbackChildFragment;
         fallbackChildFragment.effectTag |= Placement;
         primaryChildFragment.childExpirationTime = getRemainingWorkInPrimaryTree(
+          current,
           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,
+          null,
           renderExpirationTime,
         );
         // Skip the primary children, and continue working on the
         // fallback children.
-        workInProgress.memoizedState = SUSPENDED_MARKER;
+        workInProgress.memoizedState = mountSuspenseState(renderExpirationTime);
         workInProgress.child = primaryChildFragment;
         return fallbackChildFragment;
       } else {

packages/react-reconciler/src/ReactFiberHydrationContext.js

@@ -55,7 +55,7 @@ import {
   didNotFindHydratableSuspenseInstance,
 } from './ReactFiberHostConfig';
 import {enableSuspenseServerRenderer} from 'shared/ReactFeatureFlags';
-import {Never} from './ReactFiberExpirationTime';
+import {Never, NoWork} from './ReactFiberExpirationTime';
 
 // The deepest Fiber on the stack involved in a hydration context.
 // This may have been an insertion or a hydration.
@@ -231,6 +231,7 @@ function tryHydrate(fiber, nextInstance) {
         if (suspenseInstance !== null) {
           const suspenseState: SuspenseState = {
             dehydrated: suspenseInstance,
+            baseTime: NoWork,
             retryTime: Never,
           };
           fiber.memoizedState = suspenseState;

packages/react-reconciler/src/ReactFiberSuspenseComponent.js

@@ -35,6 +35,10 @@ export type SuspenseState = {|
   // here to indicate that it is dehydrated (flag) and for quick access
   // to check things like isSuspenseInstancePending.
   dehydrated: null | SuspenseInstance,
+  // Represents the work that was deprioritized when we committed the fallback.
+  // The work outside the boundary already committed at this level, so we cannot
+  // unhide the content without including it.
+  baseTime: ExpirationTime,
   // Represents the earliest expiration time we should attempt to hydrate
   // a dehydrated boundary at.
   // Never is the default for dehydrated boundaries.

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

@@ -3168,9 +3168,14 @@ describe('ReactSuspenseWithNoopRenderer', () => {
       });
 
       expect(Scheduler).toHaveYielded([
-        // First try to update the suspended tree. It's still suspended.
-        'Suspend! [C]',
+        // First try to render the high pri update. We won't try to re-render
+        // the suspended tree during this pass, because it still has unfinished
+        // updates at a lower priority.
         'Loading...',
+
+        // Now try the suspended update again. It's still suspended.
+        'Suspend! [C]',
+
         // Then complete the update to the fallback.
         'Still loading...',
       ]);
@@ -3260,4 +3265,252 @@ describe('ReactSuspenseWithNoopRenderer', () => {
       expect(root).toMatchRenderedOutput(<span prop="D" />);
     },
   );
+
+  it(
+    'after showing fallback, should not flip back to primary content until ' +
+      'the update that suspended finishes',
+    async () => {
+      const {useState, useEffect} = React;
+      const root = ReactNoop.createRoot();
+
+      let setOuterText;
+      function Parent({step}) {
+        const [text, _setText] = useState('A');
+        setOuterText = _setText;
+        return (
+          <>
+            <Text text={'Outer text: ' + text} />
+            <Text text={'Outer step: ' + step} />
+            <Suspense fallback={<Text text="Loading..." />}>
+              <Child step={step} outerText={text} />
+            </Suspense>
+          </>
+        );
+      }
+
+      let setInnerText;
+      function Child({step, outerText}) {
+        const [text, _setText] = useState('A');
+        setInnerText = _setText;
+
+        // This will log if the component commits in an inconsistent state
+        useEffect(() => {
+          if (text === outerText) {
+            Scheduler.unstable_yieldValue('Commit Child');
+          } else {
+            Scheduler.unstable_yieldValue(
+              'FIXME: Texts are inconsistent (tearing)',
+            );
+          }
+        }, [text, outerText]);
+
+        return (
+          <>
+            <AsyncText text={'Inner text: ' + text} />
+            <Text text={'Inner step: ' + step} />
+          </>
+        );
+      }
+
+      // These always update simultaneously. They must be consistent.
+      function setText(text) {
+        setOuterText(text);
+        setInnerText(text);
+      }
+
+      // Mount an initial tree. Resolve A so that it doesn't suspend.
+      await resolveText('Inner text: A');
+      await ReactNoop.act(async () => {
+        root.render(<Parent step={0} />);
+      });
+      expect(Scheduler).toHaveYielded([
+        'Outer text: A',
+        'Outer step: 0',
+        'Inner text: A',
+        'Inner step: 0',
+        'Commit Child',
+      ]);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span prop="Outer text: A" />
+          <span prop="Outer step: 0" />
+          <span prop="Inner text: A" />
+          <span prop="Inner step: 0" />
+        </>,
+      );
+
+      // Update. This causes the inner component to suspend.
+      await ReactNoop.act(async () => {
+        setText('B');
+      });
+      expect(Scheduler).toHaveYielded([
+        'Outer text: B',
+        'Outer step: 0',
+        'Suspend! [Inner text: B]',
+        'Inner step: 0',
+        'Loading...',
+      ]);
+      // Commit the placeholder
+      await advanceTimers(250);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span prop="Outer text: B" />
+          <span prop="Outer step: 0" />
+          <span hidden={true} prop="Inner text: A" />
+          <span hidden={true} prop="Inner step: 0" />
+          <span prop="Loading..." />
+        </>,
+      );
+
+      // Schedule a high pri update on the parent.
+      await ReactNoop.act(async () => {
+        ReactNoop.discreteUpdates(() => {
+          root.render(<Parent step={1} />);
+        });
+      });
+      // Only the outer part can update. The inner part should still show a
+      // fallback because we haven't finished loading B yet. Otherwise, the
+      // inner text would be inconsistent with the outer text.
+      expect(Scheduler).toHaveYielded([
+        'Outer text: B',
+        'Outer step: 1',
+        'Loading...',
+
+        'Suspend! [Inner text: B]',
+        'Inner step: 1',
+      ]);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span prop="Outer text: B" />
+          <span prop="Outer step: 1" />
+          <span hidden={true} prop="Inner text: A" />
+          <span hidden={true} prop="Inner step: 0" />
+          <span prop="Loading..." />
+        </>,
+      );
+
+      // Now finish resolving the inner text
+      await ReactNoop.act(async () => {
+        await resolveText('Inner text: B');
+      });
+      expect(Scheduler).toHaveYielded([
+        'Promise resolved [Inner text: B]',
+        'Inner text: B',
+        'Inner step: 1',
+        'Commit Child',
+      ]);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span prop="Outer text: B" />
+          <span prop="Outer step: 1" />
+          <span prop="Inner text: B" />
+          <span prop="Inner step: 1" />
+        </>,
+      );
+    },
+  );
+
+  it('a high pri update can unhide a boundary that suspended at a different level', async () => {
+    const {useState, useEffect} = React;
+    const root = ReactNoop.createRoot();
+
+    let setOuterText;
+    function Parent({step}) {
+      const [text, _setText] = useState('A');
+      setOuterText = _setText;
+      return (
+        <>
+          <Text text={'Outer: ' + text + step} />
+          <Suspense fallback={<Text text="Loading..." />}>
+            <Child step={step} outerText={text} />
+          </Suspense>
+        </>
+      );
+    }
+
+    let setInnerText;
+    function Child({step, outerText}) {
+      const [text, _setText] = useState('A');
+      setInnerText = _setText;
+
+      // This will log if the component commits in an inconsistent state
+      useEffect(() => {
+        if (text === outerText) {
+          Scheduler.unstable_yieldValue('Commit Child');
+        } else {
+          Scheduler.unstable_yieldValue(
+            'FIXME: Texts are inconsistent (tearing)',
+          );
+        }
+      }, [text, outerText]);
+
+      return (
+        <>
+          <AsyncText text={'Inner: ' + text + step} />
+        </>
+      );
+    }
+
+    // These always update simultaneously. They must be consistent.
+    function setText(text) {
+      setOuterText(text);
+      setInnerText(text);
+    }
+
+    // Mount an initial tree. Resolve A so that it doesn't suspend.
+    await resolveText('Inner: A0');
+    await ReactNoop.act(async () => {
+      root.render(<Parent step={0} />);
+    });
+    expect(Scheduler).toHaveYielded(['Outer: A0', 'Inner: A0', 'Commit Child']);
+    expect(root).toMatchRenderedOutput(
+      <>
+        <span prop="Outer: A0" />
+        <span prop="Inner: A0" />
+      </>,
+    );
+
+    // Update. This causes the inner component to suspend.
+    await ReactNoop.act(async () => {
+      setText('B');
+    });
+    expect(Scheduler).toHaveYielded([
+      'Outer: B0',
+      'Suspend! [Inner: B0]',
+      'Loading...',
+    ]);
+    // Commit the placeholder
+    await advanceTimers(250);
+    expect(root).toMatchRenderedOutput(
+      <>
+        <span prop="Outer: B0" />
+        <span hidden={true} prop="Inner: A0" />
+        <span prop="Loading..." />
+      </>,
+    );
+
+    // Schedule a high pri update on the parent. This will unblock the content.
+    await resolveText('Inner: B1');
+    await ReactNoop.act(async () => {
+      ReactNoop.discreteUpdates(() => {
+        root.render(<Parent step={1} />);
+      });
+    });
+
+    expect(Scheduler).toHaveYielded([
+      // First the outer part of the tree updates, at high pri.
+      'Outer: B1',
+      'Loading...',
+
+      // Then we retry the boundary.
+      'Inner: B1',
+      'Commit Child',
+    ]);
+    expect(root).toMatchRenderedOutput(
+      <>
+        <span prop="Outer: B1" />
+        <span prop="Inner: B1" />
+      </>,
+    );
+  });
 });