Reset hydration state after reentering (#16306)

We might reenter a hydration state, when attempting to hydrate a boundary.
We need to ensure that we reset it to not hydrating once we exit it.
Otherwise the next sibling will still be in hydration mode.

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

@@ -1057,4 +1057,43 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(container.lastChild.nodeType).toBe(8);
     expect(container.lastChild.data).toBe('unrelated comment');
   });
+
+  it('can hydrate TWO suspense boundaries', async () => {
+    let ref1 = React.createRef();
+    let ref2 = React.createRef();
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading 1...">
+            <span ref={ref1}>1</span>
+          </Suspense>
+          <Suspense fallback="Loading 2...">
+            <span ref={ref2}>2</span>
+          </Suspense>
+        </div>
+      );
+    }
+
+    // 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.
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    let span1 = container.getElementsByTagName('span')[0];
+    let span2 = container.getElementsByTagName('span')[1];
+
+    // On the client we don't have all data yet but we want to start
+    // hydrating anyway.
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App />);
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    expect(ref1.current).toBe(span1);
+    expect(ref2.current).toBe(span2);
+  });
 });

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -148,6 +148,7 @@ import {
   reenterHydrationStateFromDehydratedSuspenseInstance,
   resetHydrationState,
   tryToClaimNextHydratableInstance,
+  warnIfHydrating,
 } from './ReactFiberHydrationContext';
 import {
   adoptClassInstance,
@@ -1910,12 +1911,18 @@ function updateDehydratedSuspenseComponent(
     }
     return null;
   }
+
   if ((workInProgress.effectTag & DidCapture) !== NoEffect) {
     // Something suspended. Leave the existing children in place.
     // TODO: In non-concurrent mode, should we commit the nodes we have hydrated so far?
     workInProgress.child = null;
     return null;
   }
+
+  // We should never be hydrating at this point because it is the first pass,
+  // but after we've already committed once.
+  warnIfHydrating();
+
   if (isSuspenseInstanceFallback(suspenseInstance)) {
     // This boundary is in a permanent fallback state. In this case, we'll never
     // get an update and we'll never be able to hydrate the final content. Let's just try the

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -117,6 +117,7 @@ import {
   prepareToHydrateHostTextInstance,
   skipPastDehydratedSuspenseInstance,
   popHydrationState,
+  resetHydrationState,
 } from './ReactFiberHydrationContext';
 import {
   enableSchedulerTracing,
@@ -982,15 +983,22 @@ function completeWork(
             markSpawnedWork(Never);
           }
           skipPastDehydratedSuspenseInstance(workInProgress);
-        } else if ((workInProgress.effectTag & DidCapture) === NoEffect) {
-          // This boundary did not suspend so it's now hydrated.
-          // To handle any future suspense cases, we're going to now upgrade it
-          // to a Suspense component. We detach it from the existing current fiber.
-          current.alternate = null;
-          workInProgress.alternate = null;
-          workInProgress.tag = SuspenseComponent;
-          workInProgress.memoizedState = null;
-          workInProgress.stateNode = null;
+        } else {
+          // We should never have been in a hydration state if we didn't have a current.
+          // However, in some of those paths, we might have reentered a hydration state
+          // and then we might be inside a hydration state. In that case, we'll need to
+          // exit out of it.
+          resetHydrationState();
+          if ((workInProgress.effectTag & DidCapture) === NoEffect) {
+            // This boundary did not suspend so it's now hydrated.
+            // To handle any future suspense cases, we're going to now upgrade it
+            // to a Suspense component. We detach it from the existing current fiber.
+            current.alternate = null;
+            workInProgress.alternate = null;
+            workInProgress.tag = SuspenseComponent;
+            workInProgress.memoizedState = null;
+            workInProgress.stateNode = null;
+          }
         }
       }
       break;

packages/react-reconciler/src/ReactFiberHydrationContext.js

@@ -51,13 +51,23 @@ import {
   didNotFindHydratableSuspenseInstance,
 } from './ReactFiberHostConfig';
 import {enableSuspenseServerRenderer} from 'shared/ReactFeatureFlags';
+import warning from 'shared/warning';
 
 // The deepest Fiber on the stack involved in a hydration context.
 // This may have been an insertion or a hydration.
 let hydrationParentFiber: null | Fiber = null;
 let nextHydratableInstance: null | HydratableInstance = null;
 let isHydrating: boolean = false;
 
+function warnIfHydrating() {
+  if (__DEV__) {
+    warning(
+      !isHydrating,
+      'We should not be hydrating here. This is a bug in React. Please file a bug.',
+    );
+  }
+}
+
 function enterHydrationState(fiber: Fiber): boolean {
   if (!supportsHydration) {
     return false;
@@ -432,6 +442,7 @@ function resetHydrationState(): void {
 }
 
 export {
+  warnIfHydrating,
   enterHydrationState,
   reenterHydrationStateFromDehydratedSuspenseInstance,
   resetHydrationState,

packages/react-reconciler/src/ReactFiberUnwindWork.js

@@ -25,6 +25,7 @@ import {enableSuspenseServerRenderer} from 'shared/ReactFeatureFlags';
 
 import {popHostContainer, popHostContext} from './ReactFiberHostContext';
 import {popSuspenseContext} from './ReactFiberSuspenseContext';
+import {resetHydrationState} from './ReactFiberHydrationContext';
 import {
   isContextProvider as isLegacyContextProvider,
   popContext as popLegacyContext,
@@ -80,8 +81,12 @@ function unwindWork(
     }
     case DehydratedSuspenseComponent: {
       if (enableSuspenseServerRenderer) {
-        // TODO: popHydrationState
         popSuspenseContext(workInProgress);
+        if (workInProgress.alternate === null) {
+          // TODO: popHydrationState
+        } else {
+          resetHydrationState();
+        }
         const effectTag = workInProgress.effectTag;
         if (effectTag & ShouldCapture) {
           workInProgress.effectTag = (effectTag & ~ShouldCapture) | DidCapture;
@@ -134,7 +139,6 @@ function unwindInterruptedWork(interruptedWork: Fiber) {
       break;
     case DehydratedSuspenseComponent:
       if (enableSuspenseServerRenderer) {
-        // TODO: popHydrationState
         popSuspenseContext(interruptedWork);
       }
       break;