Check for store mutations before commit

When a store is read for the first time, or when `subscribe` or
`getSnapshot` changes, during a concurrent render, we have to check
at the end of the render phase whether the store was mutated by
an concurrent event.

In the userspace shim, we perform this check in a layout effect, and
patch up any inconsistencies by scheduling another render + commit.
However, even though we patch them up in the next render, the parent
layout effects that fire in the original render will still observe an
inconsistent tree.

In the native implementation, we can instead check for inconsistencies
right after the root is completed, before entering the commit phase. If
we do detect a mutaiton, we can discard the tree and re-render before
firing any effects. The re-render is synchronous to block further
concurrent mutations (which is also what we do to recover from tearing
bugs that result in an error). After the synchronous re-render, we can
assume the tree the tree is consistent and continue with the normal
algorithm for finishing a completed root (i.e. either suspend
or commit).

The result is that layout effects will always observe a consistent tree.

Author: acdlite

packages/react-reconciler/src/ReactFiberFlags.js

@@ -12,51 +12,53 @@ import {enableCreateEventHandleAPI} from 'shared/ReactFeatureFlags';
 export type Flags = number;
 
 // Don't change these two values. They're used by React Dev Tools.
-export const NoFlags = /*                      */ 0b00000000000000000000000;
-export const PerformedWork = /*                */ 0b00000000000000000000001;
+export const NoFlags = /*                      */ 0b000000000000000000000000;
+export const PerformedWork = /*                */ 0b000000000000000000000001;
 
 // You can change the rest (and add more).
-export const Placement = /*                    */ 0b00000000000000000000010;
-export const Update = /*                       */ 0b00000000000000000000100;
+export const Placement = /*                    */ 0b000000000000000000000010;
+export const Update = /*                       */ 0b000000000000000000000100;
 export const PlacementAndUpdate = /*           */ Placement | Update;
-export const Deletion = /*                     */ 0b00000000000000000001000;
-export const ChildDeletion = /*                */ 0b00000000000000000010000;
-export const ContentReset = /*                 */ 0b00000000000000000100000;
-export const Callback = /*                     */ 0b00000000000000001000000;
-export const DidCapture = /*                   */ 0b00000000000000010000000;
-export const Ref = /*                          */ 0b00000000000000100000000;
-export const Snapshot = /*                     */ 0b00000000000001000000000;
-export const Passive = /*                      */ 0b00000000000010000000000;
-export const Hydrating = /*                    */ 0b00000000000100000000000;
+export const Deletion = /*                     */ 0b000000000000000000001000;
+export const ChildDeletion = /*                */ 0b000000000000000000010000;
+export const ContentReset = /*                 */ 0b000000000000000000100000;
+export const Callback = /*                     */ 0b000000000000000001000000;
+export const DidCapture = /*                   */ 0b000000000000000010000000;
+export const Ref = /*                          */ 0b000000000000000100000000;
+export const Snapshot = /*                     */ 0b000000000000001000000000;
+export const Passive = /*                      */ 0b000000000000010000000000;
+export const Hydrating = /*                    */ 0b000000000000100000000000;
 export const HydratingAndUpdate = /*           */ Hydrating | Update;
-export const Visibility = /*                   */ 0b00000000001000000000000;
+export const Visibility = /*                   */ 0b000000000001000000000000;
+export const StoreConsistency = /*             */ 0b000000000010000000000000;
 
-export const LifecycleEffectMask = Passive | Update | Callback | Ref | Snapshot;
+export const LifecycleEffectMask =
+  Passive | Update | Callback | Ref | Snapshot | StoreConsistency;
 
 // Union of all commit flags (flags with the lifetime of a particular commit)
-export const HostEffectMask = /*               */ 0b00000000001111111111111;
+export const HostEffectMask = /*               */ 0b000000000011111111111111;
 
 // These are not really side effects, but we still reuse this field.
-export const Incomplete = /*                   */ 0b00000000010000000000000;
-export const ShouldCapture = /*                */ 0b00000000100000000000000;
-export const ForceUpdateForLegacySuspense = /* */ 0b00000001000000000000000;
-export const DidPropagateContext = /*          */ 0b00000010000000000000000;
-export const NeedsPropagation = /*             */ 0b00000100000000000000000;
+export const Incomplete = /*                   */ 0b000000000100000000000000;
+export const ShouldCapture = /*                */ 0b000000001000000000000000;
+export const ForceUpdateForLegacySuspense = /* */ 0b000000010000000000000000;
+export const DidPropagateContext = /*          */ 0b000000100000000000000000;
+export const NeedsPropagation = /*             */ 0b000001000000000000000000;
 
 // Static tags describe aspects of a fiber that are not specific to a render,
 // e.g. a fiber uses a passive effect (even if there are no updates on this particular render).
 // This enables us to defer more work in the unmount case,
 // since we can defer traversing the tree during layout to look for Passive effects,
 // and instead rely on the static flag as a signal that there may be cleanup work.
-export const RefStatic = /*                    */ 0b00001000000000000000000;
-export const LayoutStatic = /*                 */ 0b00010000000000000000000;
-export const PassiveStatic = /*                */ 0b00100000000000000000000;
+export const RefStatic = /*                    */ 0b000010000000000000000000;
+export const LayoutStatic = /*                 */ 0b000100000000000000000000;
+export const PassiveStatic = /*                */ 0b001000000000000000000000;
 
 // These flags allow us to traverse to fibers that have effects on mount
 // without traversing the entire tree after every commit for
 // double invoking
-export const MountLayoutDev = /*               */ 0b01000000000000000000000;
-export const MountPassiveDev = /*              */ 0b10000000000000000000000;
+export const MountLayoutDev = /*               */ 0b010000000000000000000000;
+export const MountPassiveDev = /*              */ 0b100000000000000000000000;
 
 // Groups of flags that are used in the commit phase to skip over trees that
 // don't contain effects, by checking subtreeFlags.

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

@@ -44,6 +44,7 @@ import {
   SyncLane,
   NoLanes,
   isSubsetOfLanes,
+  includesBlockingLane,
   mergeLanes,
   removeLanes,
   intersectLanes,
@@ -68,6 +69,7 @@ import {
   PassiveStatic as PassiveStaticEffect,
   StaticMask as StaticMaskEffect,
   Update as UpdateEffect,
+  StoreConsistency,
 } from './ReactFiberFlags';
 import {
   HasEffect as HookHasEffect,
@@ -166,7 +168,15 @@ type StoreInstance<T> = {|
   getSnapshot: () => T,
 |};
 
-export type FunctionComponentUpdateQueue = {|lastEffect: Effect | null|};
+type StoreConsistencyCheck<T> = {|
+  value: T,
+  getSnapshot: () => T,
+|};
+
+export type FunctionComponentUpdateQueue = {|
+  lastEffect: Effect | null,
+  stores: Array<StoreConsistencyCheck<any>> | null,
+|};
 
 type BasicStateAction<S> = (S => S) | S;
 
@@ -689,6 +699,7 @@ function updateWorkInProgressHook(): Hook {
 function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
   return {
     lastEffect: null,
+    stores: null,
   };
 }
 
@@ -1289,17 +1300,25 @@ function mountSyncExternalStore<T>(
   // don't need to set a static flag, either.
   // TODO: We can move this to the passive phase once we add a pre-commit
   // consistency check. See the next comment.
-  fiber.flags |= UpdateEffect;
+  fiber.flags |= PassiveEffect;
   pushEffect(
-    HookHasEffect | HookLayout,
+    HookHasEffect | HookPassive,
     updateStoreInstance.bind(null, fiber, inst, nextSnapshot, getSnapshot),
     undefined,
     null,
   );
-  // TODO: Unless this is a synchronous render, schedule a consistency check.
-  // Right before committing, we will walk the tree and check if any of the
-  // stores were mutated.
-  // pushConsistencyCheck(inst, getSnapshot, nextSnapshot);
+
+  // Unless we're rendering a blocking lane, schedule a consistency check. Right
+  // before committing, we will walk the tree and check if any of the stores
+  // were mutated.
+  const root: FiberRoot | null = getWorkInProgressRoot();
+  invariant(
+    root !== null,
+    'Expected a work-in-progress root. This is a bug in React. Please file an issue.',
+  );
+  if (!includesBlockingLane(root, renderLanes)) {
+    pushStoreConsistencyCheck(fiber, getSnapshot, nextSnapshot);
+  }
 
   return nextSnapshot;
 }
@@ -1348,36 +1367,69 @@ function updateSyncExternalStore<T>(
     (workInProgressHook !== null &&
       workInProgressHook.memoizedState.tag & HookHasEffect)
   ) {
-    // TODO: We can move this to the passive phase once we add a pre-commit
-    // consistency check. See the next comment.
-    fiber.flags |= UpdateEffect;
+    fiber.flags |= PassiveEffect;
     pushEffect(
-      HookHasEffect | HookLayout,
+      HookHasEffect | HookPassive,
       updateStoreInstance.bind(null, fiber, inst, nextSnapshot, getSnapshot),
       undefined,
       null,
     );
 
-    // TODO: Unless this is a synchronous render, schedule a consistency check.
+    // Unless we're rendering a blocking lane, schedule a consistency check.
     // Right before committing, we will walk the tree and check if any of the
     // stores were mutated.
-    // pushConsistencyCheck(inst, getSnapshot, nextSnapshot);
+    const root: FiberRoot | null = getWorkInProgressRoot();
+    invariant(
+      root !== null,
+      'Expected a work-in-progress root. This is a bug in React. Please file an issue.',
+    );
+    if (!includesBlockingLane(root, renderLanes)) {
+      pushStoreConsistencyCheck(fiber, getSnapshot, nextSnapshot);
+    }
   }
 
   return nextSnapshot;
 }
 
+function pushStoreConsistencyCheck<T>(
+  fiber: Fiber,
+  getSnapshot: () => T,
+  renderedSnapshot: T,
+) {
+  fiber.flags |= StoreConsistency;
+  const check: StoreConsistencyCheck<T> = {
+    getSnapshot,
+    value: renderedSnapshot,
+  };
+  let componentUpdateQueue: null | FunctionComponentUpdateQueue = (currentlyRenderingFiber.updateQueue: any);
+  if (componentUpdateQueue === null) {
+    componentUpdateQueue = createFunctionComponentUpdateQueue();
+    currentlyRenderingFiber.updateQueue = (componentUpdateQueue: any);
+    componentUpdateQueue.stores = [check];
+  } else {
+    const stores = componentUpdateQueue.stores;
+    if (stores === null) {
+      componentUpdateQueue.stores = [check];
+    } else {
+      stores.push(check);
+    }
+  }
+}
+
 function updateStoreInstance<T>(
   fiber: Fiber,
   inst: StoreInstance<T>,
   nextSnapshot: T,
   getSnapshot: () => T,
 ) {
+  // These are updated in the passive phase
   inst.value = nextSnapshot;
   inst.getSnapshot = getSnapshot;
 
-  // TODO: Move the tearing checks to an earlier, pre-commit phase so that the
-  // layout effects always observe a consistent tree.
+  // Something may have been mutated in between render and commit. This could
+  // have been in an event that fired before the passive effects, or it could
+  // have been in a layout effect. In that case, we would have used the old
+  // snapsho and getSnapshot values to bail out. We need to check one more time.
   if (checkIfSnapshotChanged(inst)) {
     // Force a re-render.
     forceStoreRerender(fiber);
@@ -1393,11 +1445,6 @@ function subscribeToStore(fiber, inst, subscribe) {
       forceStoreRerender(fiber);
     }
   };
-  // Check for changes right before subscribing. Subsequent changes will be
-  // detected in the subscription handler.
-  // TODO: Once updateStoreInstance is moved to the passive phase, we can rely
-  // on that check instead of checking again here.
-  handleStoreChange();
   // Subscribe to the store and return a clean-up function.
   return subscribe(handleStoreChange);
 }