Fix: useDeferredValue initialValue suspends forever without switching to final

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

@@ -85,6 +85,38 @@ describe('ReactDOMFizzForm', () => {
     expect(container.textContent).toEqual('Final');
   });
 
+  // @gate enableUseDeferredValueInitialArg
+  // @gate enablePostpone
+  it(
+    'if initial value postpones during hydration, it will switch to the ' +
+      'final value instead',
+    async () => {
+      function Content() {
+        const isInitial = useDeferredValue(false, true);
+        if (isInitial) {
+          React.unstable_postpone();
+        }
+        return <Text text="Final" />;
+      }
+
+      function App() {
+        return (
+          <Suspense fallback={<Text text="Loading..." />}>
+            <Content />
+          </Suspense>
+        );
+      }
+
+      const stream = await ReactDOMServer.renderToReadableStream(<App />);
+      await readIntoContainer(stream);
+      expect(container.textContent).toEqual('Loading...');
+
+      // After hydration, it's updated to the final value
+      await act(() => ReactDOMClient.hydrateRoot(container, <App />));
+      expect(container.textContent).toEqual('Final');
+    },
+  );
+
   // @gate enableUseDeferredValueInitialArg
   it(
     'useDeferredValue during hydration has higher priority than remaining ' +

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -90,6 +90,7 @@ import {
   ShouldCapture,
   ForceClientRender,
   Passive,
+  DidDefer,
 } from './ReactFiberFlags';
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 import {
@@ -257,6 +258,7 @@ import {
   renderDidSuspendDelayIfPossible,
   markSkippedUpdateLanes,
   getWorkInProgressRoot,
+  peekDeferredLane,
 } from './ReactFiberWorkLoop';
 import {enqueueConcurrentRenderForLane} from './ReactFiberConcurrentUpdates';
 import {pushCacheProvider, CacheContext} from './ReactFiberCacheComponent';
@@ -2228,9 +2230,22 @@ function shouldRemainOnFallback(
   );
 }
 
-function getRemainingWorkInPrimaryTree(current: Fiber, renderLanes: Lanes) {
-  // TODO: Should not remove render lanes that were pinged during this render
-  return removeLanes(current.childLanes, renderLanes);
+function getRemainingWorkInPrimaryTree(
+  current: Fiber | null,
+  primaryTreeDidDefer: boolean,
+  renderLanes: Lanes,
+) {
+  let remainingLanes =
+    current !== null ? removeLanes(current.childLanes, renderLanes) : NoLanes;
+  if (primaryTreeDidDefer) {
+    // A useDeferredValue hook spawned a deferred task inside the primary tree.
+    // Ensure that we retry this component at the deferred priority.
+    // TODO: We could make this a per-subtree value instead of a global one.
+    // Would need to track it on the context stack somehow, similar to what
+    // we'd have to do for resumable contexts.
+    remainingLanes = mergeLanes(remainingLanes, peekDeferredLane());
+  }
+  return remainingLanes;
 }
 
 function updateSuspenseComponent(
@@ -2259,6 +2274,11 @@ function updateSuspenseComponent(
     workInProgress.flags &= ~DidCapture;
   }
 
+  // Check if the primary children spawned a deferred task (useDeferredValue)
+  // during the first pass.
+  const didPrimaryChildrenDefer = (workInProgress.flags & DidDefer) !== NoFlags;
+  workInProgress.flags &= ~DidDefer;
+
   // OK, the next part is confusing. We're about to reconcile the Suspense
   // boundary's children. This involves some custom reconciliation logic. Two
   // main reasons this is so complicated.
@@ -2329,6 +2349,11 @@ function updateSuspenseComponent(
       const primaryChildFragment: Fiber = (workInProgress.child: any);
       primaryChildFragment.memoizedState =
         mountSuspenseOffscreenState(renderLanes);
+      primaryChildFragment.childLanes = getRemainingWorkInPrimaryTree(
+        current,
+        didPrimaryChildrenDefer,
+        renderLanes,
+      );
       workInProgress.memoizedState = SUSPENDED_MARKER;
       if (enableTransitionTracing) {
         const currentTransitions = getPendingTransitions();
@@ -2368,6 +2393,11 @@ function updateSuspenseComponent(
       const primaryChildFragment: Fiber = (workInProgress.child: any);
       primaryChildFragment.memoizedState =
         mountSuspenseOffscreenState(renderLanes);
+      primaryChildFragment.childLanes = getRemainingWorkInPrimaryTree(
+        current,
+        didPrimaryChildrenDefer,
+        renderLanes,
+      );
       workInProgress.memoizedState = SUSPENDED_MARKER;
 
       // TODO: Transition Tracing is not yet implemented for CPU Suspense.
@@ -2402,6 +2432,7 @@ function updateSuspenseComponent(
           current,
           workInProgress,
           didSuspend,
+          didPrimaryChildrenDefer,
           nextProps,
           dehydrated,
           prevState,
@@ -2464,6 +2495,7 @@ function updateSuspenseComponent(
       }
       primaryChildFragment.childLanes = getRemainingWorkInPrimaryTree(
         current,
+        didPrimaryChildrenDefer,
         renderLanes,
       );
       workInProgress.memoizedState = SUSPENDED_MARKER;
@@ -2834,6 +2866,7 @@ function updateDehydratedSuspenseComponent(
   current: Fiber,
   workInProgress: Fiber,
   didSuspend: boolean,
+  didPrimaryChildrenDefer: boolean,
   nextProps: any,
   suspenseInstance: SuspenseInstance,
   suspenseState: SuspenseState,
@@ -3063,6 +3096,11 @@ function updateDehydratedSuspenseComponent(
       const primaryChildFragment: Fiber = (workInProgress.child: any);
       primaryChildFragment.memoizedState =
         mountSuspenseOffscreenState(renderLanes);
+      primaryChildFragment.childLanes = getRemainingWorkInPrimaryTree(
+        current,
+        didPrimaryChildrenDefer,
+        renderLanes,
+      );
       workInProgress.memoizedState = SUSPENDED_MARKER;
       return fallbackChildFragment;
     }

packages/react-reconciler/src/ReactFiberFlags.js

@@ -41,6 +41,7 @@ export const StoreConsistency = /*             */ 0b0000000000000100000000000000
 // possible, because we're about to run out of bits.
 export const ScheduleRetry = StoreConsistency;
 export const ShouldSuspendCommit = Visibility;
+export const DidDefer = ContentReset;
 
 export const LifecycleEffectMask =
   Passive | Update | Callback | Ref | Snapshot | StoreConsistency;

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -127,6 +127,7 @@ import {
   Visibility,
   MountPassiveDev,
   MountLayoutDev,
+  DidDefer,
 } from './ReactFiberFlags';
 import {
   NoLanes,
@@ -714,6 +715,20 @@ export function requestDeferredLane(): Lane {
       workInProgressDeferredLane = requestTransitionLane();
     }
   }
+
+  // Mark the parent Suspense boundary so it knows to spawn the deferred lane.
+  const suspenseHandler = getSuspenseHandler();
+  if (suspenseHandler !== null) {
+    // TODO: As an optimization, we shouldn't entangle the lanes at the root; we
+    // can entangle them using the baseLanes of the Suspense boundary instead.
+    // We only need to do something special if there's no Suspense boundary.
+    suspenseHandler.flags |= DidDefer;
+  }
+
+  return workInProgressDeferredLane;
+}
+
+export function peekDeferredLane(): Lane {
   return workInProgressDeferredLane;
 }
 
@@ -1361,7 +1376,7 @@ export function performSyncWorkOnRoot(root: FiberRoot, lanes: Lanes): null {
     // The render unwound without completing the tree. This happens in special
     // cases where need to exit the current render without producing a
     // consistent tree or committing.
-    markRootSuspended(root, lanes, NoLane);
+    markRootSuspended(root, lanes, workInProgressDeferredLane);
     ensureRootIsScheduled(root);
     return null;
   }

packages/react-reconciler/src/__tests__/ReactDeferredValue-test.js

@@ -409,7 +409,7 @@ describe('ReactDeferredValue', () => {
   // @gate enableUseDeferredValueInitialArg
   it(
     'if a suspended render spawns a deferred task, we can switch to the ' +
-      'deferred task without finishing the original one',
+      'deferred task without finishing the original one (no Suspense boundary)',
     async () => {
       function App() {
         const text = useDeferredValue('Final', 'Loading...');
@@ -439,6 +439,87 @@ describe('ReactDeferredValue', () => {
     },
   );
 
+  // @gate enableUseDeferredValueInitialArg
+  it(
+    'if a suspended render spawns a deferred task, we can switch to the ' +
+      'deferred task without finishing the original one (no Suspense boundary, ' +
+      'synchronous parent update)',
+    async () => {
+      function App() {
+        const text = useDeferredValue('Final', 'Loading...');
+        return <AsyncText text={text} />;
+      }
+
+      const root = ReactNoop.createRoot();
+      // TODO: This made me realize that we don't warn if an update spawns a
+      // deferred task without being wrapped with `act`. Usually it would work
+      // anyway because the parent task has to wrapped with `act`... but not
+      // if it was flushed with `flushSync` instead.
+      await act(() => {
+        ReactNoop.flushSync(() => root.render(<App />));
+      });
+      assertLog([
+        'Suspend! [Loading...]',
+        // The initial value suspended, so we attempt the final value, which
+        // also suspends.
+        'Suspend! [Final]',
+      ]);
+      expect(root).toMatchRenderedOutput(null);
+
+      // The final value loads, so we can skip the initial value entirely.
+      await act(() => resolveText('Final'));
+      assertLog(['Final']);
+      expect(root).toMatchRenderedOutput('Final');
+
+      // When the initial value finally loads, nothing happens because we no
+      // longer need it.
+      await act(() => resolveText('Loading...'));
+      assertLog([]);
+      expect(root).toMatchRenderedOutput('Final');
+    },
+  );
+
+  // @gate enableUseDeferredValueInitialArg
+  it(
+    'if a suspended render spawns a deferred task, we can switch to the ' +
+      'deferred task without finishing the original one (Suspense boundary)',
+    async () => {
+      function App() {
+        const text = useDeferredValue('Final', 'Loading...');
+        return <AsyncText text={text} />;
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(() =>
+        root.render(
+          <Suspense fallback={<Text text="Fallback" />}>
+            <App />
+          </Suspense>,
+        ),
+      );
+      assertLog([
+        'Suspend! [Loading...]',
+        'Fallback',
+
+        // The initial value suspended, so we attempt the final value, which
+        // also suspends.
+        'Suspend! [Final]',
+      ]);
+      expect(root).toMatchRenderedOutput('Fallback');
+
+      // The final value loads, so we can skip the initial value entirely.
+      await act(() => resolveText('Final'));
+      assertLog(['Final']);
+      expect(root).toMatchRenderedOutput('Final');
+
+      // When the initial value finally loads, nothing happens because we no
+      // longer need it.
+      await act(() => resolveText('Loading...'));
+      assertLog([]);
+      expect(root).toMatchRenderedOutput('Final');
+    },
+  );
+
   // @gate enableUseDeferredValueInitialArg
   it(
     'if a suspended render spawns a deferred task that also suspends, we can ' +