Hydrate using SuspenseComponent as the parent (#22582)

* Add a failing test for Suspense hydration

* Include salazarm's changes to the test

* The hydration parent of a suspense boundary should be the boundary itself

This eventually got set when we popped back out of its children but it
doesn't start out that way.

This fixes it so that the boundary parent is always the suspense boundary.

* We now need to log errors with a suspense boundary as a parent

For now, we just log this with commentNode.parentNode as the parent for
purposes of the error message.

* Make a special getFirstHydratableChildWithinSuspenseInstance

We currently call getNextHydratableSibling but conceptually it's the child
of the boundary. It just happens to be that when we use comment nodes, we
need to call nextSibling in the DOM.

This makes this separation a bit clearer.

* Sync old fork

Co-authored-by: Dan Abramov <dan.abramov@me.com>

Author: sebmarkbage

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

@@ -160,6 +160,64 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(ref.current).toBe(span);
   });
 
+  it('can hydrate siblings of a suspended component without errors', async () => {
+    let suspend = false;
+    let resolve;
+    const promise = new Promise(resolvePromise => (resolve = resolvePromise));
+    function Child() {
+      if (suspend) {
+        throw promise;
+      } else {
+        return 'Hello';
+      }
+    }
+
+    function App() {
+      return (
+        <Suspense fallback="Loading...">
+          <Child />
+          <Suspense fallback="Loading...">
+            <div>Hello</div>
+          </Suspense>
+        </Suspense>
+      );
+    }
+
+    // First we render the final HTML. With the streaming renderer
+    // this may have suspense points on the server but here we want
+    // to test the completed HTML. Don't suspend on the server.
+    suspend = false;
+    const finalHTML = ReactDOMServer.renderToString(<App />);
+
+    const container = document.createElement('div');
+    container.innerHTML = finalHTML;
+    expect(container.textContent).toBe('HelloHello');
+
+    // On the client we don't have all data yet but we want to start
+    // hydrating anyway.
+    suspend = true;
+    ReactDOM.hydrateRoot(container, <App />);
+    expect(() => {
+      Scheduler.unstable_flushAll();
+    }).toErrorDev(
+      // TODO: This error should not be logged in this case. It's a false positive.
+      'Did not expect server HTML to contain the text node "Hello" in <div>.',
+    );
+    jest.runAllTimers();
+
+    // Expect the server-generated HTML to stay intact.
+    expect(container.textContent).toBe('HelloHello');
+
+    // Resolving the promise should continue hydration
+    suspend = false;
+    resolve();
+    await promise;
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+    // Hydration should not change anything.
+    expect(container.textContent).toBe('HelloHello');
+  });
+
   it('calls the hydration callbacks after hydration or deletion', async () => {
     let suspend = false;
     let resolve;

packages/react-dom/src/client/ReactDOMHostConfig.js

@@ -764,11 +764,23 @@ export function getNextHydratableSibling(
 }
 
 export function getFirstHydratableChild(
-  parentInstance: Container | Instance,
+  parentInstance: Instance,
 ): null | HydratableInstance {
   return getNextHydratable(parentInstance.firstChild);
 }
 
+export function getFirstHydratableChildWithinContainer(
+  parentContainer: Container,
+): null | HydratableInstance {
+  return getNextHydratable(parentContainer.firstChild);
+}
+
+export function getFirstHydratableChildWithinSuspenseInstance(
+  parentInstance: SuspenseInstance,
+): null | HydratableInstance {
+  return getNextHydratable(parentInstance.nextSibling);
+}
+
 export function hydrateInstance(
   instance: Instance,
   type: string,
@@ -917,7 +929,7 @@ export function didNotMatchHydratedTextInstance(
   }
 }
 
-export function didNotHydrateContainerInstance(
+export function didNotHydrateInstanceWithinContainer(
   parentContainer: Container,
   instance: HydratableInstance,
 ) {
@@ -932,6 +944,25 @@ export function didNotHydrateContainerInstance(
   }
 }
 
+export function didNotHydrateInstanceWithinSuspenseInstance(
+  parentInstance: SuspenseInstance,
+  instance: HydratableInstance,
+) {
+  if (__DEV__) {
+    // $FlowFixMe: Only Element or Document can be parent nodes.
+    const parentNode: Element | Document | null = parentInstance.parentNode;
+    if (parentNode !== null) {
+      if (instance.nodeType === ELEMENT_NODE) {
+        warnForDeletedHydratableElement(parentNode, (instance: any));
+      } else if (instance.nodeType === COMMENT_NODE) {
+        // TODO: warnForDeletedHydratableSuspenseBoundary
+      } else {
+        warnForDeletedHydratableText(parentNode, (instance: any));
+      }
+    }
+  }
+}
+
 export function didNotHydrateInstance(
   parentType: string,
   parentProps: Props,
@@ -949,7 +980,7 @@ export function didNotHydrateInstance(
   }
 }
 
-export function didNotFindHydratableContainerInstance(
+export function didNotFindHydratableInstanceWithinContainer(
   parentContainer: Container,
   type: string,
   props: Props,
@@ -959,7 +990,7 @@ export function didNotFindHydratableContainerInstance(
   }
 }
 
-export function didNotFindHydratableContainerTextInstance(
+export function didNotFindHydratableTextInstanceWithinContainer(
   parentContainer: Container,
   text: string,
 ) {
@@ -968,14 +999,47 @@ export function didNotFindHydratableContainerTextInstance(
   }
 }
 
-export function didNotFindHydratableContainerSuspenseInstance(
+export function didNotFindHydratableSuspenseInstanceWithinContainer(
   parentContainer: Container,
 ) {
   if (__DEV__) {
     // TODO: warnForInsertedHydratedSuspense(parentContainer);
   }
 }
 
+export function didNotFindHydratableInstanceWithinSuspenseInstance(
+  parentInstance: SuspenseInstance,
+  type: string,
+  props: Props,
+) {
+  if (__DEV__) {
+    // $FlowFixMe: Only Element or Document can be parent nodes.
+    const parentNode: Element | Document | null = parentInstance.parentNode;
+    if (parentNode !== null)
+      warnForInsertedHydratedElement(parentNode, type, props);
+  }
+}
+
+export function didNotFindHydratableTextInstanceWithinSuspenseInstance(
+  parentInstance: SuspenseInstance,
+  text: string,
+) {
+  if (__DEV__) {
+    // $FlowFixMe: Only Element or Document can be parent nodes.
+    const parentNode: Element | Document | null = parentInstance.parentNode;
+    if (parentNode !== null) warnForInsertedHydratedText(parentNode, text);
+  }
+}
+
+export function didNotFindHydratableSuspenseInstanceWithinSuspenseInstance(
+  parentInstance: SuspenseInstance,
+) {
+  if (__DEV__) {
+    // const parentNode: Element | Document | null = parentInstance.parentNode;
+    // TODO: warnForInsertedHydratedSuspense(parentNode);
+  }
+}
+
 export function didNotFindHydratableInstance(
   parentType: string,
   parentProps: Props,

packages/react-reconciler/src/ReactFiberHostConfigWithNoHydration.js

@@ -29,6 +29,8 @@ export const isSuspenseInstanceFallback = shim;
 export const registerSuspenseInstanceRetry = shim;
 export const getNextHydratableSibling = shim;
 export const getFirstHydratableChild = shim;
+export const getFirstHydratableChildWithinContainer = shim;
+export const getFirstHydratableChildWithinSuspenseInstance = shim;
 export const hydrateInstance = shim;
 export const hydrateTextInstance = shim;
 export const hydrateSuspenseInstance = shim;
@@ -40,11 +42,15 @@ export const clearSuspenseBoundaryFromContainer = shim;
 export const shouldDeleteUnhydratedTailInstances = shim;
 export const didNotMatchHydratedContainerTextInstance = shim;
 export const didNotMatchHydratedTextInstance = shim;
-export const didNotHydrateContainerInstance = shim;
+export const didNotHydrateInstanceWithinContainer = shim;
+export const didNotHydrateInstanceWithinSuspenseInstance = shim;
 export const didNotHydrateInstance = shim;
-export const didNotFindHydratableContainerInstance = shim;
-export const didNotFindHydratableContainerTextInstance = shim;
-export const didNotFindHydratableContainerSuspenseInstance = shim;
+export const didNotFindHydratableInstanceWithinContainer = shim;
+export const didNotFindHydratableTextInstanceWithinContainer = shim;
+export const didNotFindHydratableSuspenseInstanceWithinContainer = shim;
+export const didNotFindHydratableInstanceWithinSuspenseInstance = shim;
+export const didNotFindHydratableTextInstanceWithinSuspenseInstance = shim;
+export const didNotFindHydratableSuspenseInstanceWithinSuspenseInstance = shim;
 export const didNotFindHydratableInstance = shim;
 export const didNotFindHydratableTextInstance = shim;
 export const didNotFindHydratableSuspenseInstance = shim;

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

@@ -38,18 +38,24 @@ import {
   canHydrateSuspenseInstance,
   getNextHydratableSibling,
   getFirstHydratableChild,
+  getFirstHydratableChildWithinContainer,
+  getFirstHydratableChildWithinSuspenseInstance,
   hydrateInstance,
   hydrateTextInstance,
   hydrateSuspenseInstance,
   getNextHydratableInstanceAfterSuspenseInstance,
   shouldDeleteUnhydratedTailInstances,
   didNotMatchHydratedContainerTextInstance,
   didNotMatchHydratedTextInstance,
-  didNotHydrateContainerInstance,
+  didNotHydrateInstanceWithinContainer,
+  didNotHydrateInstanceWithinSuspenseInstance,
   didNotHydrateInstance,
-  didNotFindHydratableContainerInstance,
-  didNotFindHydratableContainerTextInstance,
-  didNotFindHydratableContainerSuspenseInstance,
+  didNotFindHydratableInstanceWithinContainer,
+  didNotFindHydratableTextInstanceWithinContainer,
+  didNotFindHydratableSuspenseInstanceWithinContainer,
+  didNotFindHydratableInstanceWithinSuspenseInstance,
+  didNotFindHydratableTextInstanceWithinSuspenseInstance,
+  didNotFindHydratableSuspenseInstanceWithinSuspenseInstance,
   didNotFindHydratableInstance,
   didNotFindHydratableTextInstance,
   didNotFindHydratableSuspenseInstance,
@@ -78,8 +84,10 @@ function enterHydrationState(fiber: Fiber): boolean {
     return false;
   }
 
-  const parentInstance = fiber.stateNode.containerInfo;
-  nextHydratableInstance = getFirstHydratableChild(parentInstance);
+  const parentInstance: Container = fiber.stateNode.containerInfo;
+  nextHydratableInstance = getFirstHydratableChildWithinContainer(
+    parentInstance,
+  );
   hydrationParentFiber = fiber;
   isHydrating = true;
   return true;
@@ -92,8 +100,10 @@ function reenterHydrationStateFromDehydratedSuspenseInstance(
   if (!supportsHydration) {
     return false;
   }
-  nextHydratableInstance = getNextHydratableSibling(suspenseInstance);
-  popToNextHostParent(fiber);
+  nextHydratableInstance = getFirstHydratableChildWithinSuspenseInstance(
+    suspenseInstance,
+  );
+  hydrationParentFiber = fiber;
   isHydrating = true;
   return true;
 }
@@ -105,7 +115,7 @@ function deleteHydratableInstance(
   if (__DEV__) {
     switch (returnFiber.tag) {
       case HostRoot:
-        didNotHydrateContainerInstance(
+        didNotHydrateInstanceWithinContainer(
           returnFiber.stateNode.containerInfo,
           instance,
         );
@@ -118,6 +128,14 @@ function deleteHydratableInstance(
           instance,
         );
         break;
+      case SuspenseComponent:
+        const suspenseState: SuspenseState = returnFiber.memoizedState;
+        if (suspenseState.dehydrated !== null)
+          didNotHydrateInstanceWithinSuspenseInstance(
+            suspenseState.dehydrated,
+            instance,
+          );
+        break;
     }
   }
 
@@ -144,14 +162,23 @@ function insertNonHydratedInstance(returnFiber: Fiber, fiber: Fiber) {
           case HostComponent:
             const type = fiber.type;
             const props = fiber.pendingProps;
-            didNotFindHydratableContainerInstance(parentContainer, type, props);
+            didNotFindHydratableInstanceWithinContainer(
+              parentContainer,
+              type,
+              props,
+            );
             break;
           case HostText:
             const text = fiber.pendingProps;
-            didNotFindHydratableContainerTextInstance(parentContainer, text);
+            didNotFindHydratableTextInstanceWithinContainer(
+              parentContainer,
+              text,
+            );
             break;
           case SuspenseComponent:
-            didNotFindHydratableContainerSuspenseInstance(parentContainer);
+            didNotFindHydratableSuspenseInstanceWithinContainer(
+              parentContainer,
+            );
             break;
         }
         break;
@@ -191,6 +218,35 @@ function insertNonHydratedInstance(returnFiber: Fiber, fiber: Fiber) {
         }
         break;
       }
+      case SuspenseComponent: {
+        const suspenseState: SuspenseState = returnFiber.memoizedState;
+        const parentInstance = suspenseState.dehydrated;
+        if (parentInstance !== null)
+          switch (fiber.tag) {
+            case HostComponent:
+              const type = fiber.type;
+              const props = fiber.pendingProps;
+              didNotFindHydratableInstanceWithinSuspenseInstance(
+                parentInstance,
+                type,
+                props,
+              );
+              break;
+            case HostText:
+              const text = fiber.pendingProps;
+              didNotFindHydratableTextInstanceWithinSuspenseInstance(
+                parentInstance,
+                text,
+              );
+              break;
+            case SuspenseComponent:
+              didNotFindHydratableSuspenseInstanceWithinSuspenseInstance(
+                parentInstance,
+              );
+              break;
+          }
+        break;
+      }
       default:
         return;
     }