[Hydration] Fallback to client render if server rendered extra nodes

packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js

@@ -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 () => {

packages/react-reconciler/src/ReactFiberCompleteWork.new.js

@@ -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.

packages/react-reconciler/src/ReactFiberCompleteWork.old.js

@@ -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.

packages/react-reconciler/src/ReactFiberHydrationContext.new.js

@@ -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,
 };