Mark root as already hydrated after committing (#16739)

* Mark root as already hydrated after committing

* Remove current/child check for hydration and rely on the root flag instead

packages/react-dom/src/__tests__/ReactServerRenderingHydration-test.js

@@ -502,6 +502,21 @@ describe('ReactDOMServerHydration', () => {
     expect(element.textContent).toBe('Hello world');
   });
 
+  it('does not re-enter hydration after committing the first one', () => {
+    let finalHTML = ReactDOMServer.renderToString(<div />);
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<div />);
+    Scheduler.unstable_flushAll();
+    root.render(null);
+    Scheduler.unstable_flushAll();
+    // This should not reenter hydration state and therefore not trigger hydration
+    // warnings.
+    root.render(<div />);
+    Scheduler.unstable_flushAll();
+  });
+
   it('does not invoke an event on a concurrent hydrating node until it commits', () => {
     function Sibling({text}) {
       Scheduler.unstable_yieldValue('Sibling');

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -937,14 +937,7 @@ function updateHostRoot(current, workInProgress, renderExpirationTime) {
     );
   }
   const root: FiberRoot = workInProgress.stateNode;
-  if (
-    // TODO: This is a bug because if we render null after having hydrating,
-    // we'll reenter hydration state at the next update which will then
-    // trigger hydration warnings.
-    (current === null || current.child === null) &&
-    root.hydrate &&
-    enterHydrationState(workInProgress)
-  ) {
+  if (root.hydrate && enterHydrationState(workInProgress)) {
     // If we don't have any current children this might be the first pass.
     // We always try to hydrate. If this isn't a hydration pass there won't
     // be any children to hydrate which is effectively the same thing as

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -81,6 +81,7 @@ import {
   getPublicInstance,
   supportsMutation,
   supportsPersistence,
+  supportsHydration,
   commitMount,
   commitUpdate,
   resetTextContent,
@@ -1297,6 +1298,16 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
         attachSuspenseRetryListeners(finishedWork);
         return;
       }
+      case HostRoot: {
+        const root: FiberRoot = finishedWork.stateNode;
+        if (supportsHydration) {
+          if (root.hydrate) {
+            // We've just hydrated. No need to hydrate again.
+            root.hydrate = false;
+          }
+        }
+        break;
+      }
     }
 
     commitContainer(finishedWork);
@@ -1366,6 +1377,13 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
       return;
     }
     case HostRoot: {
+      const root: FiberRoot = finishedWork.stateNode;
+      if (supportsHydration) {
+        if (root.hydrate) {
+          // We've just hydrated. No need to hydrate again.
+          root.hydrate = false;
+        }
+      }
       return;
     }
     case Profiler: {

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -662,7 +662,12 @@ function completeWork(
       if (current === null || current.child === null) {
         // If we hydrated, pop so that we can delete any remaining children
         // that weren't hydrated.
-        popHydrationState(workInProgress);
+        let wasHydrated = popHydrationState(workInProgress);
+        if (wasHydrated) {
+          // If we hydrated, then we'll need to schedule an update for
+          // the commit side-effects on the root.
+          markUpdate(workInProgress);
+        }
       }
       updateHostContainer(workInProgress);
       break;