Allow updating dehydrated root at lower priority without forcing clie…

…nt render (#24082)

* Pass children to hydration root constructor

I already made this change for the concurrent root API in #23309. This
does the same thing for the legacy API.

Doesn't change any behavior, but I will use this in the next steps.

* Add isRootDehydrated function

Currently this does nothing except read a boolean field, but I'm about
to change this logic.

Since this is accessed by React DOM, too, I put the function in a
separate module that can be deep imported. Previously, it was accessing
the FiberRoot directly. The reason it's a separate module is to break a
circular dependency between React DOM and the reconciler.

* Allow updates at lower pri without forcing client render

Currently, if a root is updated before the shell has finished hydrating
(for example, due to a top-level navigation), we immediately revert to
client rendering. This is rare because the root is expected is finish
quickly, but not exceedingly rare because the root may be suspended.

This adds support for updating the root without forcing a client render
as long as the update has lower priority than the initial hydration,
i.e. if the update is wrapped in startTransition.

To implement this, I had to do some refactoring. The main idea here is
to make it closer to how we implement hydration in Suspense boundaries:

- I moved isDehydrated from the shared FiberRoot object to the
HostRoot's state object.
- In the begin phase, I check if the root has received an by comparing
the new children to the initial children. If they are different, we
revert to client rendering, and set isDehydrated to false using a
derived state update (a la getDerivedStateFromProps).
- There are a few places where we used to set root.isDehydrated to false
as a way to force a client render. Instead, I set the ForceClientRender
flag on the root work-in-progress fiber.
- Whenever we fall back to client rendering, I log a recoverable error.

The overall code structure is almost identical to the corresponding
logic for Suspense components.

The reason this works is because if the update has lower priority than
the initial hydration, it won't be processed during the hydration
render, so the children will be the same.

We can go even further and allow updates at _higher_ priority (though
not sync) by implementing selective hydration at the root, like we do
for Suspense boundaries: interrupt the current render, attempt hydration
at slightly higher priority than the update, then continue rendering the
update. I haven't implemented this yet, but I've structured the code in
anticipation of adding this later.

* Wrap useMutableSource logic in feature flag

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,