Repository: facebook/react
Pull request: [Hydration] Fallback to client render if server rendered extra nodes

<!--
  Thanks for submitting a pull request!
  We appreciate you spending the time to work on these changes. Please provide enough information so that others can review your pull request. The three fields below are mandatory.

  Before submitting a pull request, please make sure the following is done:

  1. Fork [the repository](https://github.com/facebook/react) and create your branch from `main`.
  2. Run `yarn` in the repository root.
  3. If you've fixed a bug or added code that should be tested, add tests!
  4. Ensure the test suite passes (`yarn test`). Tip: `yarn test --watch TestName` is helpful in development.
  5. Run `yarn test --prod` to test in the production environment. It supports the same options as `yarn test`.
  6. If you need a debugger, run `yarn debug-test --watch TestName`, open `chrome://inspect`, and press "Inspect".
  7. Format your code with [prettier](https://github.com/prettier/prettier) (`yarn prettier`).
  8. Make sure your code lints (`yarn lint`). Tip: `yarn linc` to only check changed files.
  9. Run the [Flow](https://flowtype.org/) type checks (`yarn flow`).
  10. If you haven't already, complete the CLA.

  Learn more about contributing: https://reactjs.org/docs/how-to-contribute.html
-->

## Summary
In the case of extra nodes being rendered on the server thats technically a mismatch so we should follow the mismatch recovery behavior of doing a client render.

Since we throw from inside completeWork in `popHydrationState` I had to move `popTreeContext` and `popSuspenseContext` so that they happen after it that way we don't double pop when re-entering `completeWork` from `handleError`

<!--
 Explain the **motivation** for making this change. What existing problem does the pull request solve?
-->

## How did you test this change?
jest
<!--
  Demonstrate the code is solid. Example: The exact commands you ran and their output, screenshots / videos if the pull request changes the user interface.
  How exactly did you verify that your PR solves the issue you wanted to solve?
  If you leave this empty, your PR will very likely be closed.
-->


Changed files:
filename: packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js
patch: @@ -437,7 +437,7 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(container.innerHTML).toContain('<div>Sibling</div>');
   });
 
-  it('recovers when server rendered additional nodes', async () => {
+  it('recovers with client render when server rendered additional nodes at suspense root', async () => {
     const ref = React.createRef();
     function App({hasB}) {
       return (
@@ -462,15 +462,128 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(container.innerHTML).toContain('<span>B</span>');
     expect(ref.current).toBe(null);
 
-    ReactDOM.hydrateRoot(container, <App hasB={false} />);
     expect(() => {
-      Scheduler.unstable_flushAll();
+      act(() => {
+        ReactDOM.hydrateRoot(container, <App hasB={false} />);
+      });
     }).toErrorDev('Did not expect server HTML to contain a <span> in <div>');
+
     jest.runAllTimers();
 
     expect(container.innerHTML).toContain('<span>A</span>');
     expect(container.innerHTML).not.toContain('<span>B</span>');
-    expect(ref.current).toBe(span);
+
+    if (gate(flags => flags.enableClientRenderFallbackOnHydrationMismatch)) {
+      expect(ref.current).not.toBe(span);
+    } else {
+      expect(ref.current).toBe(span);
+    }
+  });
+
+  it('recovers with client render when server rendered additional nodes at suspense root after unsuspending', async () => {
+    spyOnDev(console, 'error');
+    const ref = React.createRef();
+    function App({hasB}) {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <Suspender />
+            <span ref={ref}>A</span>
+            {hasB ? <span>B</span> : null}
+          </Suspense>
+          <div>Sibling</div>
+        </div>
+      );
+    }
+
+    let shouldSuspend = false;
+    let resolve;
+    const promise = new Promise(res => {
+      resolve = () => {
+        shouldSuspend = false;
+        res();
+      };
+    });
+    function Suspender() {
+      if (shouldSuspend) {
+        throw promise;
+      }
+      return <></>;
+    }
+
+    const finalHTML = ReactDOMServer.renderToString(<App hasB={true} />);
+
+    const container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    const span = container.getElementsByTagName('span')[0];
+
+    expect(container.innerHTML).toContain('<span>A</span>');
+    expect(container.innerHTML).toContain('<span>B</span>');
+    expect(ref.current).toBe(null);
+
+    shouldSuspend = true;
+    act(() => {
+      ReactDOM.hydrateRoot(container, <App hasB={false} />);
+    });
+
+    // await expect(async () => {
+    resolve();
+    await promise;
+    Scheduler.unstable_flushAll();
+    await null;
+    jest.runAllTimers();
+    // }).toErrorDev('Did not expect server HTML to contain a <span> in <div>');
+
+    expect(container.innerHTML).toContain('<span>A</span>');
+    expect(container.innerHTML).not.toContain('<span>B</span>');
+    if (gate(flags => flags.enableClientRenderFallbackOnHydrationMismatch)) {
+      expect(ref.current).not.toBe(span);
+    } else {
+      expect(ref.current).toBe(span);
+    }
+  });
+
+  it('recovers with client render when server rendered additional nodes deep inside suspense root', async () => {
+    const ref = React.createRef();
+    function App({hasB}) {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <div>
+              <span ref={ref}>A</span>
+              {hasB ? <span>B</span> : null}
+            </div>
+          </Suspense>
+          <div>Sibling</div>
+        </div>
+      );
+    }
+
+    const finalHTML = ReactDOMServer.renderToString(<App hasB={true} />);
+
+    const container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    const span = container.getElementsByTagName('span')[0];
+
+    expect(container.innerHTML).toContain('<span>A</span>');
+    expect(container.innerHTML).toContain('<span>B</span>');
+    expect(ref.current).toBe(null);
+
+    expect(() => {
+      act(() => {
+        ReactDOM.hydrateRoot(container, <App hasB={false} />);
+      });
+    }).toErrorDev('Did not expect server HTML to contain a <span> in <div>');
+
+    expect(container.innerHTML).toContain('<span>A</span>');
+    expect(container.innerHTML).not.toContain('<span>B</span>');
+    if (gate(flags => flags.enableClientRenderFallbackOnHydrationMismatch)) {
+      expect(ref.current).not.toBe(span);
+    } else {
+      expect(ref.current).toBe(span);
+    }
   });
 
   it('calls the onDeleted hydration callback if the parent gets deleted', async () => {

filename: packages/react-reconciler/src/ReactFiberCompleteWork.new.js
patch: @@ -29,7 +29,10 @@ import type {
 import type {SuspenseContext} from './ReactFiberSuspenseContext.new';
 import type {OffscreenState} from './ReactFiberOffscreenComponent';
 import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent.new';
-import {enableSuspenseAvoidThisFallback} from 'shared/ReactFeatureFlags';
+import {
+  enableClientRenderFallbackOnHydrationMismatch,
+  enableSuspenseAvoidThisFallback,
+} from 'shared/ReactFeatureFlags';
 
 import {resetWorkInProgressVersions as resetMutableSourceWorkInProgressVersions} from './ReactMutableSource.new';
 
@@ -74,6 +77,9 @@ import {
   StaticMask,
   MutationMask,
   Passive,
+  Incomplete,
+  ShouldCapture,
+  ForceClientRender,
 } from './ReactFiberFlags';
 
 import {
@@ -120,9 +126,11 @@ import {
   prepareToHydrateHostInstance,
   prepareToHydrateHostTextInstance,
   prepareToHydrateHostSuspenseInstance,
+  warnIfUnhydratedTailNodes,
   popHydrationState,
   resetHydrationState,
   getIsHydrating,
+  hasUnhydratedTailNodes,
 } from './ReactFiberHydrationContext.new';
 import {
   enableSuspenseCallback,
@@ -1021,6 +1029,18 @@ function completeWork(
       const nextState: null | SuspenseState = workInProgress.memoizedState;
 
       if (enableSuspenseServerRenderer) {
+        if (
+          enableClientRenderFallbackOnHydrationMismatch &&
+          hasUnhydratedTailNodes() &&
+          (workInProgress.mode & ConcurrentMode) !== NoMode &&
+          (workInProgress.flags & DidCapture) === NoFlags
+        ) {
+          warnIfUnhydratedTailNodes(workInProgress);
+          resetHydrationState();
+          workInProgress.flags |=
+            ForceClientRender | Incomplete | ShouldCapture;
+          return workInProgress;
+        }
         if (nextState !== null && nextState.dehydrated !== null) {
           // We might be inside a hydration state the first time we're picking up this
           // Suspense boundary, and also after we've reentered it for further hydration.

filename: packages/react-reconciler/src/ReactFiberCompleteWork.old.js
patch: @@ -29,7 +29,10 @@ import type {
 import type {SuspenseContext} from './ReactFiberSuspenseContext.old';
 import type {OffscreenState} from './ReactFiberOffscreenComponent';
 import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent.old';
-import {enableSuspenseAvoidThisFallback} from 'shared/ReactFeatureFlags';
+import {
+  enableClientRenderFallbackOnHydrationMismatch,
+  enableSuspenseAvoidThisFallback,
+} from 'shared/ReactFeatureFlags';
 
 import {resetWorkInProgressVersions as resetMutableSourceWorkInProgressVersions} from './ReactMutableSource.old';
 
@@ -74,6 +77,9 @@ import {
   StaticMask,
   MutationMask,
   Passive,
+  Incomplete,
+  ShouldCapture,
+  ForceClientRender,
 } from './ReactFiberFlags';
 
 import {
@@ -120,9 +126,11 @@ import {
   prepareToHydrateHostInstance,
   prepareToHydrateHostTextInstance,
   prepareToHydrateHostSuspenseInstance,
+  warnIfUnhydratedTailNodes,
   popHydrationState,
   resetHydrationState,
   getIsHydrating,
+  hasUnhydratedTailNodes,
 } from './ReactFiberHydrationContext.old';
 import {
   enableSuspenseCallback,
@@ -1021,6 +1029,18 @@ function completeWork(
       const nextState: null | SuspenseState = workInProgress.memoizedState;
 
       if (enableSuspenseServerRenderer) {
+        if (
+          enableClientRenderFallbackOnHydrationMismatch &&
+          hasUnhydratedTailNodes() &&
+          (workInProgress.mode & ConcurrentMode) !== NoMode &&
+          (workInProgress.flags & DidCapture) === NoFlags
+        ) {
+          warnIfUnhydratedTailNodes(workInProgress);
+          resetHydrationState();
+          workInProgress.flags |=
+            ForceClientRender | Incomplete | ShouldCapture;
+          return workInProgress;
+        }
         if (nextState !== null && nextState.dehydrated !== null) {
           // We might be inside a hydration state the first time we're picking up this
           // Suspense boundary, and also after we've reentered it for further hydration.

filename: packages/react-reconciler/src/ReactFiberHydrationContext.new.js
patch: @@ -26,7 +26,13 @@ import {
   HostRoot,
   SuspenseComponent,
 } from './ReactWorkTags';
-import {ChildDeletion, Placement, Hydrating} from './ReactFiberFlags';
+import {
+  ChildDeletion,
+  Placement,
+  Hydrating,
+  NoFlags,
+  DidCapture,
+} from './ReactFiberFlags';
 
 import {
   createFiberFromHostInstanceForDeletion,
@@ -121,7 +127,7 @@ function reenterHydrationStateFromDehydratedSuspenseInstance(
   return true;
 }
 
-function deleteHydratableInstance(
+function warnUnhydratedInstance(
   returnFiber: Fiber,
   instance: HydratableInstance,
 ) {
@@ -151,7 +157,13 @@ function deleteHydratableInstance(
         break;
     }
   }
+}
 
+function deleteHydratableInstance(
+  returnFiber: Fiber,
+  instance: HydratableInstance,
+) {
+  warnUnhydratedInstance(returnFiber, instance);
   const childToDelete = createFiberFromHostInstanceForDeletion();
   childToDelete.stateNode = instance;
   childToDelete.return = returnFiber;
@@ -327,11 +339,16 @@ function tryHydrate(fiber, nextInstance) {
   }
 }
 
-function throwOnHydrationMismatchIfConcurrentMode(fiber: Fiber) {
-  if (
+function shouldClientRenderOnMismatch(fiber: Fiber) {
+  return (
     enableClientRenderFallbackOnHydrationMismatch &&
-    (fiber.mode & ConcurrentMode) !== NoMode
-  ) {
+    (fiber.mode & ConcurrentMode) !== NoMode &&
+    (fiber.flags & DidCapture) === NoFlags
+  );
+}
+
+function throwOnHydrationMismatchIfConcurrentMode(fiber: Fiber) {
+  if (shouldClientRenderOnMismatch(fiber)) {
     throw new Error(
       'An error occurred during hydration. The server HTML was replaced with client content',
     );
@@ -539,12 +556,18 @@ function popHydrationState(fiber: Fiber): boolean {
         !shouldSetTextContent(fiber.type, fiber.memoizedProps)))
   ) {
     let nextInstance = nextHydratableInstance;
-    while (nextInstance) {
-      deleteHydratableInstance(fiber, nextInstance);
-      nextInstance = getNextHydratableSibling(nextInstance);
+    if (nextInstance) {
+      if (shouldClientRenderOnMismatch(fiber)) {
+        warnIfUnhydratedTailNodes(fiber);
+        throwOnHydrationMismatchIfConcurrentMode(fiber);
+      } else {
+        while (nextInstance) {
+          deleteHydratableInstance(fiber, nextInstance);
+          nextInstance = getNextHydratableSibling(nextInstance);
+        }
+      }
     }
   }
-
   popToNextHostParent(fiber);
   if (fiber.tag === SuspenseComponent) {
     nextHydratableInstance = skipPastDehydratedSuspenseInstance(fiber);
@@ -556,6 +579,18 @@ function popHydrationState(fiber: Fiber): boolean {
   return true;
 }
 
+function hasUnhydratedTailNodes() {
+  return isHydrating && nextHydratableInstance !== null;
+}
+
+function warnIfUnhydratedTailNodes(fiber: Fiber) {
+  let nextInstance = nextHydratableInstance;
+  while (nextInstance) {
+    warnUnhydratedInstance(fiber, nextInstance);
+    nextInstance = getNextHydratableSibling(nextInstance);
+  }
+}
+
 function resetHydrationState(): void {
   if (!supportsHydration) {
     return;
@@ -581,4 +616,6 @@ export {
   prepareToHydrateHostTextInstance,
   prepareToHydrateHostSuspenseInstance,
   popHydrationState,
+  hasUnhydratedTailNodes,
+  warnIfUnhydratedTailNodes,
 };

filename: packages/react-reconciler/src/ReactFiberHydrationContext.old.js
patch: @@ -26,7 +26,13 @@ import {
   HostRoot,
   SuspenseComponent,
 } from './ReactWorkTags';
-import {ChildDeletion, Placement, Hydrating} from './ReactFiberFlags';
+import {
+  ChildDeletion,
+  Placement,
+  Hydrating,
+  NoFlags,
+  DidCapture,
+} from './ReactFiberFlags';
 
 import {
   createFiberFromHostInstanceForDeletion,
@@ -121,7 +127,7 @@ function reenterHydrationStateFromDehydratedSuspenseInstance(
   return true;
 }
 
-function deleteHydratableInstance(
+function warnUnhydratedInstance(
   returnFiber: Fiber,
   instance: HydratableInstance,
 ) {
@@ -151,7 +157,13 @@ function deleteHydratableInstance(
         break;
     }
   }
+}
 
+function deleteHydratableInstance(
+  returnFiber: Fiber,
+  instance: HydratableInstance,
+) {
+  warnUnhydratedInstance(returnFiber, instance);
   const childToDelete = createFiberFromHostInstanceForDeletion();
   childToDelete.stateNode = instance;
   childToDelete.return = returnFiber;
@@ -327,11 +339,16 @@ function tryHydrate(fiber, nextInstance) {
   }
 }
 
-function throwOnHydrationMismatchIfConcurrentMode(fiber: Fiber) {
-  if (
+function shouldClientRenderOnMismatch(fiber: Fiber) {
+  return (
     enableClientRenderFallbackOnHydrationMismatch &&
-    (fiber.mode & ConcurrentMode) !== NoMode
-  ) {
+    (fiber.mode & ConcurrentMode) !== NoMode &&
+    (fiber.flags & DidCapture) === NoFlags
+  );
+}
+
+function throwOnHydrationMismatchIfConcurrentMode(fiber: Fiber) {
+  if (shouldClientRenderOnMismatch(fiber)) {
     throw new Error(
       'An error occurred during hydration. The server HTML was replaced with client content',
     );
@@ -539,12 +556,18 @@ function popHydrationState(fiber: Fiber): boolean {
         !shouldSetTextContent(fiber.type, fiber.memoizedProps)))
   ) {
     let nextInstance = nextHydratableInstance;
-    while (nextInstance) {
-      deleteHydratableInstance(fiber, nextInstance);
-      nextInstance = getNextHydratableSibling(nextInstance);
+    if (nextInstance) {
+      if (shouldClientRenderOnMismatch(fiber)) {
+        warnIfUnhydratedTailNodes(fiber);
+        throwOnHydrationMismatchIfConcurrentMode(fiber);
+      } else {
+        while (nextInstance) {
+          deleteHydratableInstance(fiber, nextInstance);
+          nextInstance = getNextHydratableSibling(nextInstance);
+        }
+      }
     }
   }
-
   popToNextHostParent(fiber);
   if (fiber.tag === SuspenseComponent) {
     nextHydratableInstance = skipPastDehydratedSuspenseInstance(fiber);
@@ -556,6 +579,18 @@ function popHydrationState(fiber: Fiber): boolean {
   return true;
 }
 
+function hasUnhydratedTailNodes() {
+  return isHydrating && nextHydratableInstance !== null;
+}
+
+function warnIfUnhydratedTailNodes(fiber: Fiber) {
+  let nextInstance = nextHydratableInstance;
+  while (nextInstance) {
+    warnUnhydratedInstance(fiber, nextInstance);
+    nextInstance = getNextHydratableSibling(nextInstance);
+  }
+}
+
 function resetHydrationState(): void {
   if (!supportsHydration) {
     return;
@@ -581,4 +616,6 @@ export {
   prepareToHydrateHostTextInstance,
   prepareToHydrateHostSuspenseInstance,
   popHydrationState,
+  hasUnhydratedTailNodes,
+  warnIfUnhydratedTailNodes,
 };