useDeferredValue has higher priority than partial hydration (facebook#27550)

By default, partial hydration is given the lowest possible priority,
because until a tree is updated, the server-rendered HTML is assumed to
match the final resolved HTML.

However, this isn't completely true because a component may choose to
"upgrade" itself upon hydration. The simplest example is a component
that calls setState in a useEffect to switch to a richer implementation
of the UI. Another example is a component that doesn't have a server-
rendered implementation, so it intentionally suspends to force a client-
only render.

useDeferredValue is an example, too: the server only renders the first
pass (the initialValue) argument, and relies on the client to upgrade to
the final value.

What we should really do in these cases is emit some information into
the Fizz stream so that Fiber knows to prioritize the hydration of
certain trees. We plan to add a mechanism for this in the future.

In the meantime, though, we can at least ensure that the priority of the
upgrade task is correct once it's "discovered" during hydration. In this
case, the priority of the task spawned by useDeferredValue should have
Transition priority, not Offscreen priority.

Author: acdlite

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

@@ -9,28 +9,39 @@
 
 'use strict';
 
-import {insertNodesAndExecuteScripts} from '../test-utils/FizzTestUtils';
+import {
+  insertNodesAndExecuteScripts,
+  getVisibleChildren,
+} from '../test-utils/FizzTestUtils';
 
 // Polyfills for test environment
 global.ReadableStream =
   require('web-streams-polyfill/ponyfill/es6').ReadableStream;
 global.TextEncoder = require('util').TextEncoder;
 
 let act;
+let assertLog;
+let waitForPaint;
 let container;
 let React;
+let Scheduler;
 let ReactDOMServer;
 let ReactDOMClient;
 let useDeferredValue;
+let Suspense;
 
 describe('ReactDOMFizzForm', () => {
   beforeEach(() => {
     jest.resetModules();
     React = require('react');
+    Scheduler = require('scheduler');
     ReactDOMServer = require('react-dom/server.browser');
     ReactDOMClient = require('react-dom/client');
-    useDeferredValue = require('react').useDeferredValue;
+    useDeferredValue = React.useDeferredValue;
+    Suspense = React.Suspense;
     act = require('internal-test-utils').act;
+    assertLog = require('internal-test-utils').assertLog;
+    waitForPaint = require('internal-test-utils').waitForPaint;
     container = document.createElement('div');
     document.body.appendChild(container);
   });
@@ -54,6 +65,11 @@ describe('ReactDOMFizzForm', () => {
     insertNodesAndExecuteScripts(temp, container, null);
   }
 
+  function Text({text}) {
+    Scheduler.log(text);
+    return text;
+  }
+
   // @gate enableUseDeferredValueInitialArg
   it('returns initialValue argument, if provided', async () => {
     function App() {
@@ -68,4 +84,106 @@ describe('ReactDOMFizzForm', () => {
     await act(() => ReactDOMClient.hydrateRoot(container, <App />));
     expect(container.textContent).toEqual('Final');
   });
+
+  // @gate enableUseDeferredValueInitialArg
+  it(
+    'useDeferredValue during hydration has higher priority than remaining ' +
+      'incremental hydration',
+    async () => {
+      function B() {
+        const text = useDeferredValue('B [Final]', 'B [Initial]');
+        return <Text text={text} />;
+      }
+
+      function App() {
+        return (
+          <div>
+            <span>
+              <Text text="A" />
+            </span>
+            <Suspense fallback={<Text text="Loading..." />}>
+              <span>
+                <B />
+              </span>
+              <div>
+                <Suspense fallback={<Text text="Loading..." />}>
+                  <span id="C" ref={cRef}>
+                    <Text text="C" />
+                  </span>
+                </Suspense>
+              </div>
+            </Suspense>
+          </div>
+        );
+      }
+
+      const cRef = React.createRef();
+
+      // The server renders using the "initial" value for B.
+      const stream = await ReactDOMServer.renderToReadableStream(<App />);
+      await readIntoContainer(stream);
+      assertLog(['A', 'B [Initial]', 'C']);
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <span>A</span>
+          <span>B [Initial]</span>
+          <div>
+            <span id="C">C</span>
+          </div>
+        </div>,
+      );
+
+      const serverRenderedC = document.getElementById('C');
+
+      // On the client, we first hydrate the initial value, then upgrade
+      // to final.
+      await act(async () => {
+        ReactDOMClient.hydrateRoot(container, <App />);
+
+        // First the outermost Suspense boundary hydrates.
+        await waitForPaint(['A']);
+        expect(cRef.current).toBe(null);
+
+        // Then the next level hydrates. This level includes a useDeferredValue,
+        // so we should prioritize upgrading it before we proceed to hydrating
+        // additional levels.
+        await waitForPaint(['B [Initial]']);
+        expect(getVisibleChildren(container)).toEqual(
+          <div>
+            <span>A</span>
+            <span>B [Initial]</span>
+            <div>
+              <span id="C">C</span>
+            </div>
+          </div>,
+        );
+        expect(cRef.current).toBe(null);
+
+        // This paint should only update B. C should still be dehydrated.
+        await waitForPaint(['B [Final]']);
+        expect(getVisibleChildren(container)).toEqual(
+          <div>
+            <span>A</span>
+            <span>B [Final]</span>
+            <div>
+              <span id="C">C</span>
+            </div>
+          </div>,
+        );
+        expect(cRef.current).toBe(null);
+      });
+      // Finally we can hydrate C
+      assertLog(['C']);
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <span>A</span>
+          <span>B [Final]</span>
+          <div>
+            <span id="C">C</span>
+          </div>
+        </div>,
+      );
+      expect(cRef.current).toBe(serverRenderedC);
+    },
+  );
 });

packages/react-reconciler/src/ReactFiberLane.js

@@ -730,8 +730,10 @@ function markSpawnedDeferredLane(
   root.entanglements[spawnedLaneIndex] |=
     DeferredLane |
     // If the parent render task suspended, we must also entangle those lanes
-    // with the spawned task.
-    entangledLanes;
+    // with the spawned task, so that the deferred task includes all the same
+    // updates that the parent task did. We can exclude any lane that is not
+    // used for updates (e.g. Offscreen).
+    (entangledLanes & UpdateLanes);
 }
 
 export function markRootEntangled(root: FiberRoot, entangledLanes: Lanes) {

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -83,7 +83,10 @@ import {
   resetWorkInProgress,
 } from './ReactFiber';
 import {isRootDehydrated} from './ReactFiberShellHydration';
-import {didSuspendOrErrorWhileHydratingDEV} from './ReactFiberHydrationContext';
+import {
+  getIsHydrating,
+  didSuspendOrErrorWhileHydratingDEV,
+} from './ReactFiberHydrationContext';
 import {
   NoMode,
   ProfileMode,
@@ -690,13 +693,21 @@ export function requestDeferredLane(): Lane {
     // If there are multiple useDeferredValue hooks in the same render, the
     // tasks that they spawn should all be batched together, so they should all
     // receive the same lane.
-    if (includesSomeLane(workInProgressRootRenderLanes, OffscreenLane)) {
+
+    // Check the priority of the current render to decide the priority of the
+    // deferred task.
+
+    // OffscreenLane is used for prerendering, but we also use OffscreenLane
+    // for incremental hydration. It's given the lowest priority because the
+    // initial HTML is the same as the final UI. But useDeferredValue during
+    // hydration is an exception — we need to upgrade the UI to the final
+    // value. So if we're currently hydrating, we treat it like a transition.
+    const isPrerendering =
+      includesSomeLane(workInProgressRootRenderLanes, OffscreenLane) &&
+      !getIsHydrating();
+    if (isPrerendering) {
       // There's only one OffscreenLane, so if it contains deferred work, we
       // should just reschedule using the same lane.
-      // TODO: We also use OffscreenLane for hydration, on the basis that the
-      // initial HTML is the same as the hydrated UI, but since the deferred
-      // task will change the UI, it should be treated like an update. Use
-      // TransitionHydrationLane to trigger selective hydration.
       workInProgressDeferredLane = OffscreenLane;
     } else {
       // Everything else is spawned as a transition.