Allow updating dehydrated root at lower priority without forcing client render

packages/react-art/src/ReactART.js

@@ -69,7 +69,6 @@ class Surface extends React.Component {
     this._mountNode = createContainer(
       this._surface,
       LegacyRoot,
-      false,
       null,
       false,
       false,

packages/react-devtools-shared/src/backend/renderer.js

@@ -2704,9 +2704,14 @@ export function attach(
       // TODO: relying on this seems a bit fishy.
       const wasMounted =
         alternate.memoizedState != null &&
-        alternate.memoizedState.element != null;
+        alternate.memoizedState.element != null &&
+        // A dehydrated root is not considered mounted
+        alternate.memoizedState.isDehydrated !== true;
       const isMounted =
-        current.memoizedState != null && current.memoizedState.element != null;
+        current.memoizedState != null &&
+        current.memoizedState.element != null &&
+        // A dehydrated root is not considered mounted
+        current.memoizedState.isDehydrated !== true;
       if (!wasMounted && isMounted) {
         // Mount a new root.
         setRootPseudoKey(currentRootID, current);

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

@@ -9,6 +9,7 @@
 
 let JSDOM;
 let React;
+let startTransition;
 let ReactDOMClient;
 let Scheduler;
 let clientAct;
@@ -33,6 +34,8 @@ describe('ReactDOMFizzShellHydration', () => {
     ReactDOMFizzServer = require('react-dom/server');
     Stream = require('stream');
 
+    startTransition = React.startTransition;
+
     textCache = new Map();
 
     // Test Environment
@@ -214,7 +217,36 @@ describe('ReactDOMFizzShellHydration', () => {
     expect(container.textContent).toBe('Shell');
   });
 
-  test('updating the root before the shell hydrates forces a client render', async () => {
+  test(
+    'updating the root at lower priority than initial hydration does not ' +
+      'force a client render',
+    async () => {
+      function App() {
+        return <Text text="Initial" />;
+      }
+
+      // Server render
+      await resolveText('Initial');
+      await serverAct(async () => {
+        const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<App />);
+        pipe(writable);
+      });
+      expect(Scheduler).toHaveYielded(['Initial']);
+
+      await clientAct(async () => {
+        const root = ReactDOMClient.hydrateRoot(container, <App />);
+        // This has lower priority than the initial hydration, so the update
+        // won't be processed until after hydration finishes.
+        startTransition(() => {
+          root.render(<Text text="Updated" />);
+        });
+      });
+      expect(Scheduler).toHaveYielded(['Initial', 'Updated']);
+      expect(container.textContent).toBe('Updated');
+    },
+  );
+
+  test('updating the root while the shell is suspended forces a client render', async () => {
     function App() {
       return <AsyncText text="Shell" />;
     }
@@ -245,9 +277,9 @@ describe('ReactDOMFizzShellHydration', () => {
       root.render(<Text text="New screen" />);
     });
     expect(Scheduler).toHaveYielded([
+      'New screen',
       'This root received an early update, before anything was able ' +
         'hydrate. Switched the entire root to client rendering.',
-      'New screen',
     ]);
     expect(container.textContent).toBe('New screen');
   });

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

@@ -253,6 +253,15 @@ describe('ReactDOMRoot', () => {
     );
   });
 
+  it('callback passed to legacy hydrate() API', () => {
+    container.innerHTML = '<div>Hi</div>';
+    ReactDOM.hydrate(<div>Hi</div>, container, () => {
+      Scheduler.unstable_yieldValue('callback');
+    });
+    expect(container.textContent).toEqual('Hi');
+    expect(Scheduler).toHaveYielded(['callback']);
+  });
+
   it('warns when unmounting with legacy API (no previous content)', () => {
     const root = ReactDOMClient.createRoot(container);
     root.render(<div>Hi</div>);

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

@@ -27,6 +27,7 @@ import {
 
 import {
   createContainer,
+  createHydrationContainer,
   findHostInstanceWithNoPortals,
   updateContainer,
   flushSync,
@@ -109,34 +110,81 @@ function noopOnRecoverableError() {
 
 function legacyCreateRootFromDOMContainer(
   container: Container,
-  forceHydrate: boolean,
+  initialChildren: ReactNodeList,
+  parentComponent: ?React$Component<any, any>,
+  callback: ?Function,
+  isHydrationContainer: boolean,
 ): FiberRoot {
-  // First clear any existing content.
-  if (!forceHydrate) {
+  if (isHydrationContainer) {
+    if (typeof callback === 'function') {
+      const originalCallback = callback;
+      callback = function() {
+        const instance = getPublicRootInstance(root);
+        originalCallback.call(instance);
+      };
+    }
+
+    const root = createHydrationContainer(
+      initialChildren,
+      callback,
+      container,
+      LegacyRoot,
+      null, // hydrationCallbacks
+      false, // isStrictMode
+      false, // concurrentUpdatesByDefaultOverride,
+      '', // identifierPrefix
+      noopOnRecoverableError,
+      // TODO(luna) Support hydration later
+      null,
+    );
+    container._reactRootContainer = root;
+    markContainerAsRoot(root.current, container);
+
+    const rootContainerElement =
+      container.nodeType === COMMENT_NODE ? container.parentNode : container;
+    listenToAllSupportedEvents(rootContainerElement);
+
+    flushSync();
+    return root;
+  } else {
+    // First clear any existing content.
     let rootSibling;
     while ((rootSibling = container.lastChild)) {
       container.removeChild(rootSibling);
     }
-  }
 
-  const root = createContainer(
-    container,
-    LegacyRoot,
-    forceHydrate,
-    null, // hydrationCallbacks
-    false, // isStrictMode
-    false, // concurrentUpdatesByDefaultOverride,
-    '', // identifierPrefix
-    noopOnRecoverableError, // onRecoverableError
-    null, // transitionCallbacks
-  );
-  markContainerAsRoot(root.current, container);
+    if (typeof callback === 'function') {
+      const originalCallback = callback;
+      callback = function() {
+        const instance = getPublicRootInstance(root);
+        originalCallback.call(instance);
+      };
+    }
+
+    const root = createContainer(
+      container,
+      LegacyRoot,
+      null, // hydrationCallbacks
+      false, // isStrictMode
+      false, // concurrentUpdatesByDefaultOverride,
+      '', // identifierPrefix
+      noopOnRecoverableError, // onRecoverableError
+      null, // transitionCallbacks
+    );
+    container._reactRootContainer = root;
+    markContainerAsRoot(root.current, container);
+
+    const rootContainerElement =
+      container.nodeType === COMMENT_NODE ? container.parentNode : container;
+    listenToAllSupportedEvents(rootContainerElement);
 
-  const rootContainerElement =
-    container.nodeType === COMMENT_NODE ? container.parentNode : container;
-  listenToAllSupportedEvents(rootContainerElement);
+    // Initial mount should not be batched.
+    flushSync(() => {
+      updateContainer(initialChildren, root, parentComponent, callback);
+    });
 
-  return root;
+    return root;
+  }
 }
 
 function warnOnInvalidCallback(callback: mixed, callerName: string): void {
@@ -164,39 +212,30 @@ function legacyRenderSubtreeIntoContainer(
     warnOnInvalidCallback(callback === undefined ? null : callback, 'render');
   }
 
-  let root = container._reactRootContainer;
-  let fiberRoot: FiberRoot;
-  if (!root) {
+  const maybeRoot = container._reactRootContainer;
+  let root: FiberRoot;
+  if (!maybeRoot) {
     // Initial mount
-    root = container._reactRootContainer = legacyCreateRootFromDOMContainer(
+    root = legacyCreateRootFromDOMContainer(
       container,
+      children,
+      parentComponent,
+      callback,
       forceHydrate,
     );
-    fiberRoot = root;
-    if (typeof callback === 'function') {
-      const originalCallback = callback;
-      callback = function() {
-        const instance = getPublicRootInstance(fiberRoot);
-        originalCallback.call(instance);
-      };
-    }
-    // Initial mount should not be batched.
-    flushSync(() => {
-      updateContainer(children, fiberRoot, parentComponent, callback);
-    });
   } else {
-    fiberRoot = root;
+    root = maybeRoot;
     if (typeof callback === 'function') {
       const originalCallback = callback;
       callback = function() {
-        const instance = getPublicRootInstance(fiberRoot);
+        const instance = getPublicRootInstance(root);
         originalCallback.call(instance);
       };
     }
     // Update
-    updateContainer(children, fiberRoot, parentComponent, callback);
+    updateContainer(children, root, parentComponent, callback);
   }
-  return getPublicRootInstance(fiberRoot);
+  return getPublicRootInstance(root);
 }
 
 export function findDOMNode(

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

@@ -224,7 +224,6 @@ export function createRoot(
   const root = createContainer(
     container,
     ConcurrentRoot,
-    false,
     null,
     isStrictMode,
     concurrentUpdatesByDefaultOverride,
@@ -303,6 +302,7 @@ export function hydrateRoot(
 
   const root = createHydrationContainer(
     initialChildren,
+    null,
     container,
     ConcurrentRoot,
     hydrationCallbacks,

packages/react-dom/src/events/ReactDOMEventListener.js

@@ -53,6 +53,7 @@ import {
   setCurrentUpdatePriority,
 } from 'react-reconciler/src/ReactEventPriorities';
 import ReactSharedInternals from 'shared/ReactSharedInternals';
+import {isRootDehydrated} from 'react-reconciler/src/ReactFiberShellHydration';
 
 const {ReactCurrentBatchConfig} = ReactSharedInternals;
 
@@ -386,7 +387,7 @@ export function findInstanceBlockingEvent(
         targetInst = null;
       } else if (tag === HostRoot) {
         const root: FiberRoot = nearestMounted.stateNode;
-        if (root.isDehydrated) {
+        if (isRootDehydrated(root)) {
           // If this happens during a replay something went wrong and it might block
           // the whole system.
           return getContainerFromFiber(nearestMounted);

packages/react-dom/src/events/ReactDOMEventReplaying.js

@@ -39,6 +39,7 @@ import {
 } from '../client/ReactDOMComponentTree';
 import {HostRoot, SuspenseComponent} from 'react-reconciler/src/ReactWorkTags';
 import {isHigherEventPriority} from 'react-reconciler/src/ReactEventPriorities';
+import {isRootDehydrated} from 'react-reconciler/src/ReactFiberShellHydration';
 
 let _attemptSynchronousHydration: (fiber: Object) => void;
 
@@ -414,7 +415,7 @@ function attemptExplicitHydrationTarget(
         }
       } else if (tag === HostRoot) {
         const root: FiberRoot = nearestMounted.stateNode;
-        if (root.isDehydrated) {
+        if (isRootDehydrated(root)) {
           queuedTarget.blockedOn = getContainerFromFiber(nearestMounted);
           // We don't currently have a way to increase the priority of
           // a root other than sync.

packages/react-native-renderer/src/ReactFabric.js

@@ -215,7 +215,6 @@ function render(
     root = createContainer(
       containerTag,
       concurrentRoot ? ConcurrentRoot : LegacyRoot,
-      false,
       null,
       false,
       null,

packages/react-native-renderer/src/ReactNativeRenderer.js

@@ -211,7 +211,6 @@ function render(
     root = createContainer(
       containerTag,
       LegacyRoot,
-      false,
       null,
       false,
       null,

packages/react-noop-renderer/src/createReactNoop.js

@@ -974,7 +974,6 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
         root = NoopRenderer.createContainer(
           container,
           tag,
-          false,
           null,
           null,
           false,
@@ -996,7 +995,6 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
       const fiberRoot = NoopRenderer.createContainer(
         container,
         ConcurrentRoot,
-        false,
         null,
         null,
         false,
@@ -1029,7 +1027,6 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
       const fiberRoot = NoopRenderer.createContainer(
         container,
         LegacyRoot,
-        false,
         null,
         null,
         false,