Track entangled lanes separately from update lane (facebook#27505)

A small refactor to how the lane entanglement mechanism works. We can
now distinguish between the lane that "spawned" a render task (i.e. a
new update) versus the lanes that it's entangled with. Both the update
lane and the entangled lanes will be included while rendering, but by
keeping them separate, we don't lose the original priority.

In practical terms, this means we can now entangle a low priority update
with a higher priority lane while rendering at the lower priority.

To do this, lanes that are entangled at the root are now tracked using
the same variable that we use to track the "base lanes" when revealing a
previously hidden tree — conceptually, they are the same thing. I also
renamed this variable (from subtreeLanes to entangledRenderLanes) to
better reflect how it's used.

My primary motivation is related to useDeferredValue, which I'll address
in a later PR.

Author: acdlite

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

@@ -16,6 +16,8 @@ let ReactDOMClient;
 let Scheduler;
 let act;
 let waitForAll;
+let waitFor;
+let waitForMicrotasks;
 let assertLog;
 
 const setUntrackedInputValue = Object.getOwnPropertyDescriptor(
@@ -36,6 +38,8 @@ describe('ReactDOMFiberAsync', () => {
 
     const InternalTestUtils = require('internal-test-utils');
     waitForAll = InternalTestUtils.waitForAll;
+    waitFor = InternalTestUtils.waitFor;
+    waitForMicrotasks = InternalTestUtils.waitForMicrotasks;
     assertLog = InternalTestUtils.assertLog;
 
     document.body.appendChild(container);
@@ -653,52 +657,94 @@ describe('ReactDOMFiberAsync', () => {
     });
   });
 
-  it('transition lane in popState should yield if it suspends', async () => {
-    const never = {then() {}};
-    let _setText;
+  it('transition lane in popState should be allowed to suspend', async () => {
+    let resolvePromise;
+    const promise = new Promise(res => {
+      resolvePromise = res;
+    });
+
+    function Text({text}) {
+      Scheduler.log(text);
+      return text;
+    }
 
     function App() {
-      const [shouldSuspend, setShouldSuspend] = React.useState(false);
-      const [text, setText] = React.useState('0');
-      _setText = setText;
-      if (shouldSuspend) {
-        Scheduler.log('Suspend!');
-        throw never;
-      }
-      function onPopstate() {
-        React.startTransition(() => {
-          setShouldSuspend(val => !val);
-        });
+      const [pathname, setPathname] = React.useState('/path/a');
+
+      if (pathname !== '/path/a') {
+        try {
+          React.use(promise);
+        } catch (e) {
+          Scheduler.log(`Suspend! [${pathname}]`);
+          throw e;
+        }
       }
+
       React.useEffect(() => {
+        function onPopstate() {
+          React.startTransition(() => {
+            setPathname('/path/b');
+          });
+        }
         window.addEventListener('popstate', onPopstate);
         return () => window.removeEventListener('popstate', onPopstate);
       }, []);
-      Scheduler.log(`Child:${shouldSuspend}/${text}`);
-      return text;
+
+      return (
+        <>
+          <Text text="Before" />
+          <div>
+            <Text text={pathname} />
+          </div>
+          <Text text="After" />
+        </>
+      );
     }
 
     const root = ReactDOMClient.createRoot(container);
     await act(async () => {
       root.render(<App />);
     });
-    assertLog(['Child:false/0']);
+    assertLog(['Before', '/path/a', 'After']);
 
-    await act(() => {
+    const div = container.getElementsByTagName('div')[0];
+    expect(div.textContent).toBe('/path/a');
+
+    // Simulate a popstate event
+    await act(async () => {
       const popStateEvent = new Event('popstate');
+
+      // Simulate a popstate event
       window.event = popStateEvent;
       window.dispatchEvent(popStateEvent);
-      queueMicrotask(() => {
-        window.event = undefined;
-      });
+      await waitForMicrotasks();
+      window.event = undefined;
+
+      // The transition lane should have been attempted synchronously (in
+      // a microtask)
+      assertLog(['Suspend! [/path/b]']);
+      // Because it suspended, it remains on the current path
+      expect(div.textContent).toBe('/path/a');
     });
-    assertLog(['Suspend!']);
+    assertLog(['Suspend! [/path/b]']);
 
     await act(async () => {
-      _setText('1');
-    });
-    assertLog(['Child:false/1', 'Suspend!']);
+      resolvePromise();
+
+      // Since the transition previously suspended, there's no need for this
+      // transition to be rendered synchronously on susbequent attempts; if we
+      // fail to commit synchronously the first time, the scroll restoration
+      // state won't be restored anyway.
+      //
+      // Yield in between each child to prove that it's concurrent.
+      await waitForMicrotasks();
+      assertLog([]);
 
-    root.unmount();
+      await waitFor(['Before']);
+      await waitFor(['/path/b']);
+      await waitFor(['After']);
+    });
+    assertLog([]);
+    expect(div.textContent).toBe('/path/b');
   });
 });

packages/react-reconciler/src/ReactFiberHiddenContext.js

@@ -13,7 +13,10 @@ import type {Lanes} from './ReactFiberLane';
 
 import {createCursor, push, pop} from './ReactFiberStack';
 
-import {getRenderLanes, setRenderLanes} from './ReactFiberWorkLoop';
+import {
+  getEntangledRenderLanes,
+  setEntangledRenderLanes,
+} from './ReactFiberWorkLoop';
 import {NoLanes, mergeLanes} from './ReactFiberLane';
 
 // TODO: Remove `renderLanes` context in favor of hidden context
@@ -29,26 +32,28 @@ type HiddenContext = {
 // InvisibleParentContext that is currently managed by SuspenseContext.
 export const currentTreeHiddenStackCursor: StackCursor<HiddenContext | null> =
   createCursor(null);
-export const prevRenderLanesStackCursor: StackCursor<Lanes> =
+export const prevEntangledRenderLanesCursor: StackCursor<Lanes> =
   createCursor(NoLanes);
 
 export function pushHiddenContext(fiber: Fiber, context: HiddenContext): void {
-  const prevRenderLanes = getRenderLanes();
-  push(prevRenderLanesStackCursor, prevRenderLanes, fiber);
+  const prevEntangledRenderLanes = getEntangledRenderLanes();
+  push(prevEntangledRenderLanesCursor, prevEntangledRenderLanes, fiber);
   push(currentTreeHiddenStackCursor, context, fiber);
 
   // When rendering a subtree that's currently hidden, we must include all
   // lanes that would have rendered if the hidden subtree hadn't been deferred.
   // That is, in order to reveal content from hidden -> visible, we must commit
   // all the updates that we skipped when we originally hid the tree.
-  setRenderLanes(mergeLanes(prevRenderLanes, context.baseLanes));
+  setEntangledRenderLanes(
+    mergeLanes(prevEntangledRenderLanes, context.baseLanes),
+  );
 }
 
 export function reuseHiddenContextOnStack(fiber: Fiber): void {
   // This subtree is not currently hidden, so we don't need to add any lanes
   // to the render lanes. But we still need to push something to avoid a
   // context mismatch. Reuse the existing context on the stack.
-  push(prevRenderLanesStackCursor, getRenderLanes(), fiber);
+  push(prevEntangledRenderLanesCursor, getEntangledRenderLanes(), fiber);
   push(
     currentTreeHiddenStackCursor,
     currentTreeHiddenStackCursor.current,
@@ -58,10 +63,10 @@ export function reuseHiddenContextOnStack(fiber: Fiber): void {
 
 export function popHiddenContext(fiber: Fiber): void {
   // Restore the previous render lanes from the stack
-  setRenderLanes(prevRenderLanesStackCursor.current);
+  setEntangledRenderLanes(prevEntangledRenderLanesCursor.current);
 
   pop(currentTreeHiddenStackCursor, fiber);
-  pop(prevRenderLanesStackCursor, fiber);
+  pop(prevEntangledRenderLanesCursor, fiber);
 }
 
 export function isCurrentTreeHidden(): boolean {

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1525,7 +1525,8 @@ function mountSyncExternalStore<T>(
       );
     }
 
-    if (!includesBlockingLane(root, renderLanes)) {
+    const rootRenderLanes = getWorkInProgressRootRenderLanes();
+    if (!includesBlockingLane(root, rootRenderLanes)) {
       pushStoreConsistencyCheck(fiber, getSnapshot, nextSnapshot);
     }
   }

packages/react-reconciler/src/ReactFiberLane.js

@@ -39,6 +39,7 @@ export const NoLane: Lane = /*                          */ 0b0000000000000000000
 
 export const SyncHydrationLane: Lane = /*               */ 0b0000000000000000000000000000001;
 export const SyncLane: Lane = /*                        */ 0b0000000000000000000000000000010;
+export const SyncLaneIndex: number = 1;
 
 export const InputContinuousHydrationLane: Lane = /*    */ 0b0000000000000000000000000000100;
 export const InputContinuousLane: Lane = /*             */ 0b0000000000000000000000000001000;
@@ -274,17 +275,23 @@ export function getNextLanes(root: FiberRoot, wipLanes: Lanes): Lanes {
     }
   }
 
+  return nextLanes;
+}
+
+export function getEntangledLanes(root: FiberRoot, renderLanes: Lanes): Lanes {
+  let entangledLanes = renderLanes;
+
   if (
     allowConcurrentByDefault &&
     (root.current.mode & ConcurrentUpdatesByDefaultMode) !== NoMode
   ) {
     // Do nothing, use the lanes as they were assigned.
-  } else if ((nextLanes & InputContinuousLane) !== NoLanes) {
+  } else if ((entangledLanes & InputContinuousLane) !== NoLanes) {
     // When updates are sync by default, we entangle continuous priority updates
     // and default updates, so they render in the same batch. The only reason
     // they use separate lanes is because continuous updates should interrupt
     // transitions, but default updates should not.
-    nextLanes |= pendingLanes & DefaultLane;
+    entangledLanes |= entangledLanes & DefaultLane;
   }
 
   // Check for entangled lanes and add them to the batch.
@@ -309,21 +316,21 @@ export function getNextLanes(root: FiberRoot, wipLanes: Lanes): Lanes {
   // For those exceptions where entanglement is semantically important,
   // we should ensure that there is no partial work at the
   // time we apply the entanglement.
-  const entangledLanes = root.entangledLanes;
-  if (entangledLanes !== NoLanes) {
+  const allEntangledLanes = root.entangledLanes;
+  if (allEntangledLanes !== NoLanes) {
     const entanglements = root.entanglements;
-    let lanes = nextLanes & entangledLanes;
+    let lanes = entangledLanes & allEntangledLanes;
     while (lanes > 0) {
       const index = pickArbitraryLaneIndex(lanes);
       const lane = 1 << index;
 
-      nextLanes |= entanglements[index];
+      entangledLanes |= entanglements[index];
 
       lanes &= ~lane;
     }
   }
 
-  return nextLanes;
+  return entangledLanes;
 }
 
 function computeExpirationTime(lane: Lane, currentTime: number) {
@@ -404,6 +411,7 @@ export function markStarvedLanesAsExpired(
   // Iterate through the pending lanes and check if we've reached their
   // expiration time. If so, we'll assume the update is being starved and mark
   // it as expired to force it to finish.
+  // TODO: We should be able to replace this with upgradePendingLanesToSync
   //
   // We exclude retry lanes because those must always be time sliced, in order
   // to unwrap uncached promises.
@@ -708,6 +716,34 @@ export function markRootEntangled(root: FiberRoot, entangledLanes: Lanes) {
   }
 }
 
+export function upgradePendingLaneToSync(root: FiberRoot, lane: Lane) {
+  // Since we're upgrading the priority of the given lane, there is now pending
+  // sync work.
+  root.pendingLanes |= SyncLane;
+
+  // Entangle the sync lane with the lane we're upgrading. This means SyncLane
+  // will not be allowed to finish without also finishing the given lane.
+  root.entangledLanes |= SyncLane;
+  root.entanglements[SyncLaneIndex] |= lane;
+}
+
+export function upgradePendingLanesToSync(
+  root: FiberRoot,
+  lanesToUpgrade: Lanes,
+) {
+  // Same as upgradePendingLaneToSync but accepts multiple lanes, so it's a
+  // bit slower.
+  root.pendingLanes |= SyncLane;
+  root.entangledLanes |= SyncLane;
+  let lanes = lanesToUpgrade;
+  while (lanes) {
+    const index = pickArbitraryLaneIndex(lanes);
+    const lane = 1 << index;
+    root.entanglements[SyncLaneIndex] |= lane;
+    lanes &= ~lane;
+  }
+}
+
 export function markHiddenUpdate(
   root: FiberRoot,
   update: ConcurrentUpdate,

packages/react-reconciler/src/ReactFiberRootScheduler.js

@@ -20,8 +20,7 @@ import {
   getNextLanes,
   includesSyncLane,
   markStarvedLanesAsExpired,
-  markRootEntangled,
-  mergeLanes,
+  upgradePendingLaneToSync,
   claimNextTransitionLane,
 } from './ReactFiberLane';
 import {
@@ -250,7 +249,10 @@ function processRootScheduleInMicrotask() {
       currentEventTransitionLane !== NoLane &&
       shouldAttemptEagerTransition()
     ) {
-      markRootEntangled(root, mergeLanes(currentEventTransitionLane, SyncLane));
+      // A transition was scheduled during an event, but we're going to try to
+      // render it synchronously anyway. We do this during a popstate event to
+      // preserve the scroll position of the previous page.
+      upgradePendingLaneToSync(root, currentEventTransitionLane);
     }
 
     const nextLanes = scheduleTaskForRootDuringMicrotask(root, currentTime);