From b0a8a3e04139eecf3ad46266abdab65e5abe60c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Markb=C3=A5ge?= <sema@fb.com>
Date: Tue, 10 Sep 2019 20:02:02 -0700
Subject: [PATCH] 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
---
 .../ReactServerRenderingHydration-test.js      | 15 +++++++++++++++
 .../src/ReactFiberBeginWork.js                 |  9 +--------
 .../src/ReactFiberCommitWork.js                | 18 ++++++++++++++++++
 .../src/ReactFiberCompleteWork.js              |  7 ++++++-
 4 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactServerRenderingHydration-test.js b/packages/react-dom/src/__tests__/ReactServerRenderingHydration-test.js
index 819855bbfcee9..82cc1443729f4 100644
--- a/packages/react-dom/src/__tests__/ReactServerRenderingHydration-test.js
+++ b/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');
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index 945191c44547c..eb4234d4aef07 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/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
diff --git a/packages/react-reconciler/src/ReactFiberCommitWork.js b/packages/react-reconciler/src/ReactFiberCommitWork.js
index 7ac9433cfcb97..6a59ce8607272 100644
--- a/packages/react-reconciler/src/ReactFiberCommitWork.js
+++ b/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: {
diff --git a/packages/react-reconciler/src/ReactFiberCompleteWork.js b/packages/react-reconciler/src/ReactFiberCompleteWork.js
index 1ba603a4b7b21..77dd42dfebd75 100644
--- a/packages/react-reconciler/src/ReactFiberCompleteWork.js
+++ b/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;