useDeferredValue switches to final if initial suspends

If a parent render spawns a deferred task with useDeferredValue, but
the parent render suspends, we should not wait for the parent render to
complete before attempting to render the final value.

The reason is that the initialValue argument to useDeferredValue is
meant to represent an immediate preview of the final UI. If we can't
render it "immediately", we might as well skip it and go straight to
the "real" value.

This is an improvement over how a userspace implementation of
useDeferredValue would work, because a userspace implementation would
have to wait for the parent task to commit (useEffect) before spawning
the deferred task, creating a waterfall.

Author: acdlite

packages/react-reconciler/src/ReactFiberHooks.js

@@ -61,16 +61,17 @@ import {
   NoLane,
   SyncLane,
   OffscreenLane,
+  DeferredLane,
   NoLanes,
   isSubsetOfLanes,
   includesBlockingLane,
   includesOnlyNonUrgentLanes,
-  claimNextTransitionLane,
   mergeLanes,
   removeLanes,
   intersectLanes,
   isTransitionLane,
   markRootEntangled,
+  includesSomeLane,
 } from './ReactFiberLane';
 import {
   ContinuousEventPriority,
@@ -101,6 +102,7 @@ import {
   getWorkInProgressRootRenderLanes,
   scheduleUpdateOnFiber,
   requestUpdateLane,
+  requestDeferredLane,
   markSkippedUpdateLanes,
   isInvalidExecutionContextForEventFunction,
 } from './ReactFiberWorkLoop';
@@ -2665,16 +2667,21 @@ function rerenderDeferredValue<T>(value: T, initialValue?: T): T {
 }
 
 function mountDeferredValueImpl<T>(hook: Hook, value: T, initialValue?: T): T {
-  if (enableUseDeferredValueInitialArg && initialValue !== undefined) {
+  if (
+    enableUseDeferredValueInitialArg &&
     // When `initialValue` is provided, we defer the initial render even if the
     // current render is not synchronous.
-    // TODO: However, to avoid waterfalls, we should not defer if this render
-    // was itself spawned by an earlier useDeferredValue. Plan is to add a
-    // Deferred lane to track this.
+    initialValue !== undefined &&
+    // However, to avoid waterfalls, we do not defer if this render
+    // was itself spawned by an earlier useDeferredValue. Check if DeferredLane
+    // is part of the render lanes.
+    !includesSomeLane(renderLanes, DeferredLane)
+  ) {
+    // Render with the initial value
     hook.memoizedState = initialValue;
 
-    // Schedule a deferred render
-    const deferredLane = claimNextTransitionLane();
+    // Schedule a deferred render to switch to the final value.
+    const deferredLane = requestDeferredLane();
     currentlyRenderingFiber.lanes = mergeLanes(
       currentlyRenderingFiber.lanes,
       deferredLane,
@@ -2710,7 +2717,7 @@ function updateDeferredValueImpl<T>(
 
     if (!is(value, prevValue)) {
       // Schedule a deferred render
-      const deferredLane = claimNextTransitionLane();
+      const deferredLane = requestDeferredLane();
       currentlyRenderingFiber.lanes = mergeLanes(
         currentlyRenderingFiber.lanes,
         deferredLane,

packages/react-reconciler/src/ReactFiberLane.js

@@ -619,7 +619,11 @@ export function markRootUpdated(root: FiberRoot, updateLane: Lane) {
   }
 }
 
-export function markRootSuspended(root: FiberRoot, suspendedLanes: Lanes) {
+export function markRootSuspended(
+  root: FiberRoot,
+  suspendedLanes: Lanes,
+  spawnedLane: Lane,
+) {
   root.suspendedLanes |= suspendedLanes;
   root.pingedLanes &= ~suspendedLanes;
 
@@ -634,13 +638,21 @@ export function markRootSuspended(root: FiberRoot, suspendedLanes: Lanes) {
 
     lanes &= ~lane;
   }
+
+  if (spawnedLane !== NoLane) {
+    markSpawnedDeferredLane(root, spawnedLane, suspendedLanes);
+  }
 }
 
 export function markRootPinged(root: FiberRoot, pingedLanes: Lanes) {
   root.pingedLanes |= root.suspendedLanes & pingedLanes;
 }
 
-export function markRootFinished(root: FiberRoot, remainingLanes: Lanes) {
+export function markRootFinished(
+  root: FiberRoot,
+  remainingLanes: Lanes,
+  spawnedLane: Lane,
+) {
   const noLongerPendingLanes = root.pendingLanes & ~remainingLanes;
 
   root.pendingLanes = remainingLanes;
@@ -686,6 +698,37 @@ export function markRootFinished(root: FiberRoot, remainingLanes: Lanes) {
 
     lanes &= ~lane;
   }
+
+  if (spawnedLane !== NoLane) {
+    markSpawnedDeferredLane(
+      root,
+      spawnedLane,
+      // This render finished successfully without suspending, so we don't need
+      // to entangle the spawned task with the parent task.
+      NoLanes,
+    );
+  }
+}
+
+function markSpawnedDeferredLane(
+  root: FiberRoot,
+  spawnedLane: Lane,
+  entangledLanes: Lanes,
+) {
+  // This render spawned a deferred task. Mark it as pending.
+  root.pendingLanes |= spawnedLane;
+  root.suspendedLanes &= ~spawnedLane;
+
+  // Entangle the spawned lane with the DeferredLane bit so that we know it
+  // was the result of another render. This lets us avoid a useDeferredValue
+  // waterfall — only the first level will defer.
+  const spawnedLaneIndex = laneToIndex(spawnedLane);
+  root.entangledLanes |= spawnedLane;
+  root.entanglements[spawnedLaneIndex] |=
+    DeferredLane |
+    // If the parent render task suspended, we must also entangle those lanes
+    // with the spawned task.
+    entangledLanes;
 }
 
 export function markRootEntangled(root: FiberRoot, entangledLanes: Lanes) {

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -366,6 +366,8 @@ let workInProgressRootInterleavedUpdatedLanes: Lanes = NoLanes;
 let workInProgressRootRenderPhaseUpdatedLanes: Lanes = NoLanes;
 // Lanes that were pinged (in an interleaved event) during this render.
 let workInProgressRootPingedLanes: Lanes = NoLanes;
+// If this lane scheduled deferred work, this is the lane of the deferred task.
+let workInProgressDeferredLane: Lane = NoLane;
 // Errors that are thrown during the render phase.
 let workInProgressRootConcurrentErrors: Array<CapturedValue<mixed>> | null =
   null;
@@ -683,6 +685,27 @@ function requestRetryLane(fiber: Fiber) {
   return claimNextRetryLane();
 }
 
+export function requestDeferredLane(): Lane {
+  if (workInProgressDeferredLane === NoLane) {
+    // 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 (includesOnlyRetries(workInProgressRootRenderLanes)) {
+      // Retries are slightly lower priority than transitions, so if Retry task
+      // spawns a deferred task, the deferred task is also considered a Retry.
+      workInProgressDeferredLane = claimNextRetryLane();
+    } else if (includesSomeLane(workInProgressRootRenderLanes, OffscreenLane)) {
+      // There's only one OffscreenLane, so if it contains deferred work, we
+      // should just reschedule using the same lane.
+      workInProgressDeferredLane = OffscreenLane;
+    } else {
+      // Everything else is spawned as a transition.
+      workInProgressDeferredLane = requestTransitionLane();
+    }
+  }
+  return workInProgressDeferredLane;
+}
+
 export function scheduleUpdateOnFiber(
   root: FiberRoot,
   fiber: Fiber,
@@ -712,7 +735,11 @@ export function scheduleUpdateOnFiber(
     // The incoming update might unblock the current render. Interrupt the
     // current attempt and restart from the top.
     prepareFreshStack(root, NoLanes);
-    markRootSuspended(root, workInProgressRootRenderLanes);
+    markRootSuspended(
+      root,
+      workInProgressRootRenderLanes,
+      workInProgressDeferredLane,
+    );
   }
 
   // Mark that the root has a pending update.
@@ -792,7 +819,11 @@ export function scheduleUpdateOnFiber(
         // effect of interrupting the current render and switching to the update.
         // TODO: Make sure this doesn't override pings that happen while we've
         // already started rendering.
-        markRootSuspended(root, workInProgressRootRenderLanes);
+        markRootSuspended(
+          root,
+          workInProgressRootRenderLanes,
+          workInProgressDeferredLane,
+        );
       }
     }
 
@@ -903,7 +934,7 @@ export function performConcurrentWorkOnRoot(
         // 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);
+        markRootSuspended(root, lanes, NoLane);
       } else {
         // The render completed.
 
@@ -947,7 +978,7 @@ export function performConcurrentWorkOnRoot(
         if (exitStatus === RootFatalErrored) {
           const fatalError = workInProgressRootFatalError;
           prepareFreshStack(root, NoLanes);
-          markRootSuspended(root, lanes);
+          markRootSuspended(root, lanes, NoLane);
           ensureRootIsScheduled(root);
           throw fatalError;
         }
@@ -1074,7 +1105,7 @@ function finishConcurrentRender(
         // This is a transition, so we should exit without committing a
         // placeholder and without scheduling a timeout. Delay indefinitely
         // until we receive more data.
-        markRootSuspended(root, lanes);
+        markRootSuspended(root, lanes, workInProgressDeferredLane);
         return;
       }
       // Commit the placeholder.
@@ -1096,6 +1127,7 @@ function finishConcurrentRender(
       root,
       workInProgressRootRecoverableErrors,
       workInProgressTransitions,
+      workInProgressDeferredLane,
     );
   } else {
     if (
@@ -1109,7 +1141,7 @@ function finishConcurrentRender(
 
       // Don't bother with a very short suspense time.
       if (msUntilTimeout > 10) {
-        markRootSuspended(root, lanes);
+        markRootSuspended(root, lanes, workInProgressDeferredLane);
 
         const nextLanes = getNextLanes(root, NoLanes);
         if (nextLanes !== NoLanes) {
@@ -1131,6 +1163,7 @@ function finishConcurrentRender(
             workInProgressRootRecoverableErrors,
             workInProgressTransitions,
             lanes,
+            workInProgressDeferredLane,
           ),
           msUntilTimeout,
         );
@@ -1143,6 +1176,7 @@ function finishConcurrentRender(
       workInProgressRootRecoverableErrors,
       workInProgressTransitions,
       lanes,
+      workInProgressDeferredLane,
     );
   }
 }
@@ -1153,6 +1187,7 @@ function commitRootWhenReady(
   recoverableErrors: Array<CapturedValue<mixed>> | null,
   transitions: Array<Transition> | null,
   lanes: Lanes,
+  spawnedLane: Lane,
 ) {
   // TODO: Combine retry throttling with Suspensey commits. Right now they run
   // one after the other.
@@ -1180,13 +1215,13 @@ function commitRootWhenReady(
       root.cancelPendingCommit = schedulePendingCommit(
         commitRoot.bind(null, root, recoverableErrors, transitions),
       );
-      markRootSuspended(root, lanes);
+      markRootSuspended(root, lanes, spawnedLane);
       return;
     }
   }
 
   // Otherwise, commit immediately.
-  commitRoot(root, recoverableErrors, transitions);
+  commitRoot(root, recoverableErrors, transitions, spawnedLane);
 }
 
 function isRenderConsistentWithExternalStores(finishedWork: Fiber): boolean {
@@ -1242,7 +1277,11 @@ function isRenderConsistentWithExternalStores(finishedWork: Fiber): boolean {
   return true;
 }
 
-function markRootSuspended(root: FiberRoot, suspendedLanes: Lanes) {
+function markRootSuspended(
+  root: FiberRoot,
+  suspendedLanes: Lanes,
+  spawnedLane: Lane,
+) {
   // When suspending, we should always exclude lanes that were pinged or (more
   // rarely, since we try to avoid it) updated during the render phase.
   // TODO: Lol maybe there's a better way to factor this besides this
@@ -1252,7 +1291,7 @@ function markRootSuspended(root: FiberRoot, suspendedLanes: Lanes) {
     suspendedLanes,
     workInProgressRootInterleavedUpdatedLanes,
   );
-  markRootSuspended_dontCallThisOneDirectly(root, suspendedLanes);
+  markRootSuspended_dontCallThisOneDirectly(root, suspendedLanes, spawnedLane);
 }
 
 // This is the entry point for synchronous tasks that don't go
@@ -1302,7 +1341,7 @@ export function performSyncWorkOnRoot(root: FiberRoot, lanes: Lanes): null {
   if (exitStatus === RootFatalErrored) {
     const fatalError = workInProgressRootFatalError;
     prepareFreshStack(root, NoLanes);
-    markRootSuspended(root, lanes);
+    markRootSuspended(root, lanes, NoLane);
     ensureRootIsScheduled(root);
     throw fatalError;
   }
@@ -1311,7 +1350,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);
+    markRootSuspended(root, lanes, NoLane);
     ensureRootIsScheduled(root);
     return null;
   }
@@ -1325,6 +1364,7 @@ export function performSyncWorkOnRoot(root: FiberRoot, lanes: Lanes): null {
     root,
     workInProgressRootRecoverableErrors,
     workInProgressTransitions,
+    workInProgressDeferredLane,
   );
 
   // Before exiting, make sure there's a callback scheduled for the next
@@ -1537,6 +1577,7 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
   workInProgressRootInterleavedUpdatedLanes = NoLanes;
   workInProgressRootRenderPhaseUpdatedLanes = NoLanes;
   workInProgressRootPingedLanes = NoLanes;
+  workInProgressDeferredLane = NoLane;
   workInProgressRootConcurrentErrors = null;
   workInProgressRootRecoverableErrors = null;
 
@@ -1808,9 +1849,9 @@ export function renderDidSuspendDelayIfPossible(): void {
   // Check if there are updates that we skipped tree that might have unblocked
   // this render.
   if (
-    workInProgressRoot !== null &&
     (includesNonIdleWork(workInProgressRootSkippedLanes) ||
-      includesNonIdleWork(workInProgressRootInterleavedUpdatedLanes))
+      includesNonIdleWork(workInProgressRootInterleavedUpdatedLanes)) &&
+    workInProgressRoot !== null
   ) {
     // Mark the current render as suspended so that we switch to working on
     // the updates that were skipped. Usually we only suspend at the end of
@@ -1821,8 +1862,11 @@ export function renderDidSuspendDelayIfPossible(): void {
     // pinged or updated while we were rendering.
     // TODO: Consider unwinding immediately, using the
     // SuspendedOnHydration mechanism.
-    // $FlowFixMe[incompatible-call] need null check workInProgressRoot
-    markRootSuspended(workInProgressRoot, workInProgressRootRenderLanes);
+    markRootSuspended(
+      workInProgressRoot,
+      workInProgressRootRenderLanes,
+      workInProgressDeferredLane,
+    );
   }
 }
 
@@ -2592,6 +2636,7 @@ function commitRoot(
   root: FiberRoot,
   recoverableErrors: null | Array<CapturedValue<mixed>>,
   transitions: Array<Transition> | null,
+  spawnedLane: Lane,
 ) {
   // TODO: This no longer makes any sense. We already wrap the mutation and
   // layout phases. Should be able to remove.
@@ -2606,6 +2651,7 @@ function commitRoot(
       recoverableErrors,
       transitions,
       previousUpdateLanePriority,
+      spawnedLane,
     );
   } finally {
     ReactCurrentBatchConfig.transition = prevTransition;
@@ -2620,6 +2666,7 @@ function commitRootImpl(
   recoverableErrors: null | Array<CapturedValue<mixed>>,
   transitions: Array<Transition> | null,
   renderPriorityLevel: EventPriority,
+  spawnedLane: Lane,
 ) {
   do {
     // `flushPassiveEffects` will call `flushSyncUpdateQueue` at the end, which
@@ -2696,7 +2743,7 @@ function commitRootImpl(
   const concurrentlyUpdatedLanes = getConcurrentlyUpdatedLanes();
   remainingLanes = mergeLanes(remainingLanes, concurrentlyUpdatedLanes);
 
-  markRootFinished(root, remainingLanes);
+  markRootFinished(root, remainingLanes, spawnedLane);
 
   if (root === workInProgressRoot) {
     // We can reset these now that they are finished.