useMutableSource: "Entangle" instead of expiring (#18889)

* useMutableSource: "Entangle" instead of expiring A lane is said to be entangled with another when it's not allowed to render in a batch that does not also include the other lane. This commit implements entanglement for `useMutableSource`. If a source is mutated in between when it's read in the render phase, but before it's subscribed to in the commit phase, we must account for whether the same source has pending mutations elsewhere. The old subscriptions must not be allowed to re-render without also including the new subscription (and vice versa), to prevent tearing. In the old reconciler, we did this by synchronously flushing all the pending subscription updates. This works, but isn't ideal. The new reconciler can entangle the updates without de-opting to sync. In the future, we plan to use this same mechanism for other features, like skipping over intermediate useTransition states. * Use clz instead of ctrz to pick an arbitrary lane Should be slightly faster since most engines have built-in support.

Author: acdlite

packages/react-reconciler/src/ReactFiberHooks.new.js

@@ -32,7 +32,7 @@ import {
  isSubsetOfLanes,
  mergeLanes,
  removeLanes,
- markRootExpired,
+ markRootEntangled,
  markRootMutableRead,
 } from './ReactFiberLane';
 import {readContext} from './ReactFiberNewContext.new';
@@ -999,12 +999,11 @@ function useMutableSource<Source, Snapshot>(
  markRootMutableRead(root, lane);
 
  // If the source mutated between render and now,
- // there may be state updates already scheduled from the old getSnapshot.
- // Those updates should not commit without this value.
- // There is no mechanism currently to associate these updates though,
- // so for now we fall back to synchronously flushing all pending updates.
- // TODO: This should entangle the lanes instead of expiring everything.
- markRootExpired(root, root.mutableReadLanes);
+ // there may be state updates already scheduled from the old source.
+ // Entangle the updates so that they render in the same batch.
+ // TODO: I think we need to entangle even if the snapshot matches,
+ // because there could have been an update to a different hook.
+ markRootEntangled(root, root.mutableReadLanes);
  }
  }
  }, [getSnapshot, source, subscribe]);

packages/react-reconciler/src/ReactFiberLane.js

@@ -365,6 +365,37 @@ export function getNextLanes(root: FiberRoot, wipLanes: Lanes): Lanes {
  }
  }
 
+ // Check for entangled lanes and add them to the batch.
+ //
+ // A lane is said to be entangled with another when it's not allowed to render
+ // in a batch that does not also include the other lane. Typically we do this
+ // when multiple updates have the same source, and we only want to respond to
+ // the most recent event from that source.
+ //
+ // Note that we apply entanglements *after* checking for partial work above.
+ // This means that if a lane is entangled during an interleaved event while
+ // it's already rendering, we won't interrupt it. This is intentional, since
+ // entanglement is usually "best effort": we'll try our best to render the
+ // lanes in the same batch, but it's not worth throwing out partially
+ // completed work in order to do it.
+ //
+ // For those exceptions where entanglement is semantically important, like
+ // useMutableSource, we should ensure that there is no partial work at the
+ // time we apply the entanglement.
+ const entangledLanes = root.entangledLanes;
+ if (entangledLanes !== NoLanes) {
+ const entanglements = root.entanglements;
+ let lanes = nextLanes & entangledLanes;
+ while (lanes > 0) {
+ const index = pickArbitraryLaneIndex(lanes);
+ const lane = 1 << index;
+
+ nextLanes |= entanglements[index];
+
+ lanes &= ~lane;
+ }
+ }
+
  return nextLanes;
 }
 
@@ -401,7 +432,7 @@ export function markStarvedLanesAsExpired(
  // it as expired to force it to finish.
  let lanes = pendingLanes;
  while (lanes > 0) {
- const index = ctrz(lanes);
+ const index = pickArbitraryLaneIndex(lanes);
  const lane = 1 << index;
 
  const expirationTime = expirationTimes[index];
@@ -588,6 +619,10 @@ export function pickArbitraryLane(lanes: Lanes): Lane {
  return getLowestPriorityLane(lanes);
 }
 
+function pickArbitraryLaneIndex(lanes: Lane | Lanes) {
+ return 31 - clz32(lanes);
+}
+
 export function includesSomeLane(a: Lanes | Lane, b: Lanes | Lane) {
  return (a & b) !== NoLanes;
 }
@@ -647,7 +682,7 @@ export function markRootSuspended(root: FiberRoot, suspendedLanes: Lanes) {
  const expirationTimes = root.expirationTimes;
  let lanes = suspendedLanes;
  while (lanes > 0) {
- const index = ctrz(lanes);
+ const index = pickArbitraryLaneIndex(lanes);
  const lane = 1 << index;
 
  expirationTimes[index] = NoTimestamp;
@@ -692,10 +727,12 @@ export function markRootFinished(root: FiberRoot, remainingLanes: Lanes) {
  root.expiredLanes &= remainingLanes;
  root.mutableReadLanes &= remainingLanes;
 
+ root.entangledLanes &= remainingLanes;
+
  const expirationTimes = root.expirationTimes;
  let lanes = noLongerPendingLanes;
  while (lanes > 0) {
- const index = ctrz(lanes);
+ const index = pickArbitraryLaneIndex(lanes);
  const lane = 1 << index;
 
  // Clear the expiration time
@@ -705,6 +742,21 @@ export function markRootFinished(root: FiberRoot, remainingLanes: Lanes) {
  }
 }
 
+export function markRootEntangled(root: FiberRoot, entangledLanes: Lanes) {
+ root.entangledLanes |= entangledLanes;
+
+ const entanglements = root.entanglements;
+ let lanes = entangledLanes;
+ while (lanes > 0) {
+ const index = pickArbitraryLaneIndex(lanes);
+ const lane = 1 << index;
+
+ entanglements[index] |= entangledLanes;
+
+ lanes &= ~lane;
+ }
+}
+
 export function getBumpedLaneForHydration(
  root: FiberRoot,
  renderLanes: Lanes,
@@ -777,14 +829,3 @@ function clz32Fallback(lanes: Lanes | Lane) {
  }
  return (31 - ((log(lanes) / LN2) | 0)) | 0;
 }
-
-// Count trailing zeros. Only used on lanes, so assume input is an integer.
-function ctrz(lanes: Lanes | Lane) {
- let bits = lanes;
- bits |= bits << 16;
- bits |= bits << 8;
- bits |= bits << 4;
- bits |= bits << 2;
- bits |= bits << 1;
- return 32 - clz32(~bits);
-}

packages/react-reconciler/src/ReactFiberRoot.new.js

@@ -49,6 +49,9 @@ function FiberRootNode(containerInfo, tag, hydrate) {
 
  this.finishedLanes = NoLanes;
 
+ this.entangledLanes = NoLanes;
+ this.entanglements = createLaneMap(NoLanes);
+
  if (enableSchedulerTracing) {
  this.interactionThreadID = unstable_getThreadID();
  this.memoizedInteractions = new Set();

packages/react-reconciler/src/ReactFiberWorkLoop.new.js

@@ -673,6 +673,15 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
  currentEventWipLanes = NoLanes;
  currentEventPendingLanes = NoLanes;
 
+ invariant(
+ (executionContext & (RenderContext | CommitContext)) === NoContext,
+ 'Should not already be working.',
+ );
+
+ // Flush any pending passive effects before deciding which lanes to work on,
+ // in case they schedule additional work.
+ flushPassiveEffects();
+
  // Determine the next expiration time to work on, using the fields stored
  // on the root.
  let lanes = getNextLanes(
@@ -697,12 +706,6 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
  }
 
  const originalCallbackNode = root.callbackNode;
- invariant(
- (executionContext & (RenderContext | CommitContext)) === NoContext,
- 'Should not already be working.',
- );
-
- flushPassiveEffects();
 
  let exitStatus = renderRootConcurrent(root, lanes);
 

packages/react-reconciler/src/ReactInternalTypes.js

@@ -262,6 +262,9 @@ type BaseFiberRootProperties = {|
  mutableReadLanes: Lanes,
 
  finishedLanes: Lanes,
+
+ entangledLanes: Lanes,
+ entanglements: LaneMap<Lanes>,
 |};
 
 // The following attributes are only used by interaction tracing builds.

packages/react-reconciler/src/__tests__/useMutableSource-test.internal.js

@@ -1481,21 +1481,39 @@ describe('useMutableSource', () => {
  source.valueB = '3';
  },
  );
- });
-
- expect(Scheduler).toHaveYielded([
- // The partial render completes
- 'Child: 2',
- 'Commit: 2, 2',
 
- // Then we start rendering the low priority mutation
- 'Parent: 3',
+ expect(Scheduler).toFlushAndYieldThrough([
+ // The partial render completes
+ 'Child: 2',
+ 'Commit: 2, 2',
+ ]);
 
- // Eventually the child corrects itself, because of the check that
- // occurs when re-subscribing.
- 'Child: 3',
- 'Commit: 3, 3',
- ]);
+ // Now there are two pending mutations at different priorities. But they
+ // both read the same verion of the mutable source, so we must render
+ // them simultaneously.
+ //
+ if (gate(flags => flags.new)) {
+ // In the new reconciler, we can do this with entanglement: when the
+ // high priority render starts, we'll also include the low pri work.
+ expect(Scheduler).toFlushAndYieldThrough([
+ 'Parent: 3',
+ // Demonstrates that we can yield here
+ ]);
+ expect(Scheduler).toFlushAndYield([
+ // Now finish the rest of the update
+ 'Child: 3',
+ 'Commit: 3, 3',
+ ]);
+ } else {
+ // In the old reconciler, we don't have an entanglement mechanism. The
+ // best we can do is synchronously flush both updates.
+ expect(Scheduler).toFlushAndYield([
+ 'Parent: 3',
+ 'Child: 3',
+ 'Commit: 3, 3',
+ ]);
+ }
+ });
  });
 
  // @gate experimental