useMutableSource: bugfix for new getSnapshot with mutation (#18297)

packages/react-reconciler/src/ReactFiberHooks.js

@@ -31,6 +31,7 @@ import type {
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 import {enableUseEventAPI} from 'shared/ReactFeatureFlags';
 
+import {markRootExpiredAtTime} from './ReactFiberRoot';
 import {NoWork, Sync} from './ReactFiberExpirationTime';
 import {readContext} from './ReactFiberNewContext';
 import {createDeprecatedResponderListener} from './ReactFiberDeprecatedEvents';
@@ -74,7 +75,7 @@ import {
   getCurrentPriorityLevel,
 } from './SchedulerWithReactIntegration';
 import {
-  getPendingExpirationTime,
+  getLastPendingExpirationTime,
   getWorkInProgressVersion,
   markSourceAsDirty,
   setPendingExpirationTime,
@@ -886,6 +887,7 @@ function rerenderReducer<S, I, A>(
 type MutableSourceMemoizedState<Source, Snapshot> = {|
   refs: {
     getSnapshot: MutableSourceGetSnapshotFn<Source, Snapshot>,
+    setSnapshot: Snapshot => void,
   },
   source: MutableSource<any>,
   subscribe: MutableSourceSubscribeFn<Source, Snapshot>,
@@ -914,16 +916,14 @@ function readFromUnsubcribedMutableSource<Source, Snapshot>(
     isSafeToReadFromSource = currentRenderVersion === version;
   } else {
     // If there's no version, then we should fallback to checking the update time.
-    const pendingExpirationTime = getPendingExpirationTime(root);
+    const pendingExpirationTime = getLastPendingExpirationTime(root);
 
     if (pendingExpirationTime === NoWork) {
       isSafeToReadFromSource = true;
     } else {
       // If the source has pending updates, we can use the current render's expiration
       // time to determine if it's safe to read again from the source.
-      isSafeToReadFromSource =
-        pendingExpirationTime === NoWork ||
-        pendingExpirationTime >= renderExpirationTime;
+      isSafeToReadFromSource = pendingExpirationTime >= renderExpirationTime;
     }
 
     if (isSafeToReadFromSource) {
@@ -972,7 +972,8 @@ function useMutableSource<Source, Snapshot>(
 
   const dispatcher = ReactCurrentDispatcher.current;
 
-  const [currentSnapshot, setSnapshot] = dispatcher.useState(() =>
+  // eslint-disable-next-line prefer-const
+  let [currentSnapshot, setSnapshot] = dispatcher.useState(() =>
     readFromUnsubcribedMutableSource(root, source, getSnapshot),
   );
   let snapshot = currentSnapshot;
@@ -998,19 +999,51 @@ function useMutableSource<Source, Snapshot>(
     subscribe,
   }: MutableSourceMemoizedState<Source, Snapshot>);
 
-  // Sync the values needed by our subscribe function after each commit.
+  // Sync the values needed by our subscription handler after each commit.
   dispatcher.useEffect(() => {
     refs.getSnapshot = getSnapshot;
-  }, [getSnapshot]);
 
-  // If we got a new source or subscribe function,
-  // we'll need to subscribe in a passive effect,
-  // and also check for any changes that fire between render and subscribe.
+    // Normally the dispatch function for a state hook never changes,
+    // but this hook recreates the queue in certain cases  to avoid updates from stale sources.
+    // handleChange() below needs to reference the dispatch function without re-subscribing,
+    // so we use a ref to ensure that it always has the latest version.
+    refs.setSnapshot = setSnapshot;
+
+    // Check for a possible change between when we last rendered now.
+    const maybeNewVersion = getVersion(source._source);
+    if (!is(version, maybeNewVersion)) {
+      const maybeNewSnapshot = getSnapshot(source._source);
+      if (!is(snapshot, maybeNewSnapshot)) {
+        setSnapshot(maybeNewSnapshot);
+
+        const currentTime = requestCurrentTimeForUpdate();
+        const suspenseConfig = requestCurrentSuspenseConfig();
+        const expirationTime = computeExpirationForFiber(
+          currentTime,
+          fiber,
+          suspenseConfig,
+        );
+        setPendingExpirationTime(root, expirationTime);
+
+        // 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: Improve this later.
+        markRootExpiredAtTime(root, getLastPendingExpirationTime(root));
+      }
+    }
+  }, [getSnapshot, source, subscribe]);
+
+  // If we got a new source or subscribe function, re-subscribe in a passive effect.
   dispatcher.useEffect(() => {
     const handleChange = () => {
       const latestGetSnapshot = refs.getSnapshot;
+      const latestSetSnapshot = refs.setSnapshot;
+
       try {
-        setSnapshot(latestGetSnapshot(source._source));
+        latestSetSnapshot(latestGetSnapshot(source._source));
 
         // Record a pending mutable source update with the same expiration time.
         const currentTime = requestCurrentTimeForUpdate();
@@ -1027,9 +1060,11 @@ function useMutableSource<Source, Snapshot>(
         // e.g. it might try to read from a part of the store that no longer exists.
         // In this case we should still schedule an update with React.
         // Worst case the selector will throw again and then an error boundary will handle it.
-        setSnapshot(() => {
-          throw error;
-        });
+        latestSetSnapshot(
+          (() => {
+            throw error;
+          }: any),
+        );
       }
     };
 
@@ -1042,15 +1077,6 @@ function useMutableSource<Source, Snapshot>(
       }
     }
 
-    // Check for a possible change between when we last rendered and when we just subscribed.
-    const maybeNewVersion = getVersion(source._source);
-    if (!is(version, maybeNewVersion)) {
-      const maybeNewSnapshot = getSnapshot(source._source);
-      if (!is(snapshot, maybeNewSnapshot)) {
-        setSnapshot(maybeNewSnapshot);
-      }
-    }
-
     return unsubscribe;
   }, [source, subscribe]);
 
@@ -1066,10 +1092,26 @@ function useMutableSource<Source, Snapshot>(
   // In both cases, we need to throw away pending udpates (since they are no longer relevant)
   // and treat reading from the source as we do in the mount case.
   if (
+    !is(prevGetSnapshot, getSnapshot) ||
     !is(prevSource, source) ||
-    !is(prevSubscribe, subscribe) ||
-    !is(prevGetSnapshot, getSnapshot)
+    !is(prevSubscribe, subscribe)
   ) {
+    // Create a new queue and setState method,
+    // So if there are interleaved updates, they get pushed to the older queue.
+    // When this becomes current, the previous queue and dispatch method will be discarded,
+    // including any interleaving updates that occur.
+    const newQueue = {
+      pending: null,
+      dispatch: null,
+      lastRenderedReducer: basicStateReducer,
+      lastRenderedState: snapshot,
+    };
+    newQueue.dispatch = setSnapshot = (dispatchAction.bind(
+      null,
+      currentlyRenderingFiber,
+      newQueue,
+    ): any);
+    stateHook.queue = newQueue;
     stateHook.baseQueue = null;
     snapshot = readFromUnsubcribedMutableSource(root, source, getSnapshot);
     stateHook.memoizedState = stateHook.baseState = snapshot;
@@ -1087,6 +1129,7 @@ function mountMutableSource<Source, Snapshot>(
   hook.memoizedState = ({
     refs: {
       getSnapshot,
+      setSnapshot: (null: any),
     },
     source,
     subscribe,

packages/react-reconciler/src/ReactFiberRoot.js

@@ -79,7 +79,7 @@ type BaseFiberRootProperties = {|
   lastExpiredTime: ExpirationTime,
   // Used by useMutableSource hook to avoid tearing within this root
   // when external, mutable sources are read from during render.
-  mutableSourcePendingUpdateTime: ExpirationTime,
+  mutableSourceLastPendingUpdateTime: ExpirationTime,
 |};
 
 // The following attributes are only used by interaction tracing builds.
@@ -130,7 +130,8 @@ function FiberRootNode(containerInfo, tag, hydrate) {
   this.nextKnownPendingLevel = NoWork;
   this.lastPingedTime = NoWork;
   this.lastExpiredTime = NoWork;
-  this.mutableSourcePendingUpdateTime = NoWork;
+  this.mutableSourceFirstPendingUpdateTime = NoWork;
+  this.mutableSourceLastPendingUpdateTime = NoWork;
 
   if (enableSchedulerTracing) {
     this.interactionThreadID = unstable_getThreadID();

packages/react-reconciler/src/ReactMutableSource.js

@@ -30,20 +30,28 @@ export function clearPendingUpdates(
   root: FiberRoot,
   expirationTime: ExpirationTime,
 ): void {
-  if (root.mutableSourcePendingUpdateTime <= expirationTime) {
-    root.mutableSourcePendingUpdateTime = NoWork;
+  if (expirationTime <= root.mutableSourceLastPendingUpdateTime) {
+    // All updates for this source have been processed.
+    root.mutableSourceLastPendingUpdateTime = NoWork;
   }
 }
 
-export function getPendingExpirationTime(root: FiberRoot): ExpirationTime {
-  return root.mutableSourcePendingUpdateTime;
+export function getLastPendingExpirationTime(root: FiberRoot): ExpirationTime {
+  return root.mutableSourceLastPendingUpdateTime;
 }
 
 export function setPendingExpirationTime(
   root: FiberRoot,
   expirationTime: ExpirationTime,
 ): void {
-  root.mutableSourcePendingUpdateTime = expirationTime;
+  const mutableSourceLastPendingUpdateTime =
+    root.mutableSourceLastPendingUpdateTime;
+  if (
+    mutableSourceLastPendingUpdateTime === NoWork ||
+    expirationTime < mutableSourceLastPendingUpdateTime
+  ) {
+    root.mutableSourceLastPendingUpdateTime = expirationTime;
+  }
 }
 
 export function markSourceAsDirty(mutableSource: MutableSource<any>): void {

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

@@ -43,8 +43,8 @@ describe('useMutableSource', () => {
     const callbacksA = [];
     const callbacksB = [];
     let revision = 0;
-    let valueA = 'a:one';
-    let valueB = 'b:one';
+    let valueA = initialValueA;
+    let valueB = initialValueB;
 
     const subscribeHelper = (callbacks, callback) => {
       if (callbacks.indexOf(callback) < 0) {
@@ -302,7 +302,7 @@ describe('useMutableSource', () => {
           />,
           () => Scheduler.unstable_yieldValue('Sync effect'),
         );
-        expect(Scheduler).toFlushAndYieldThrough(['only:a-one', 'Sync effect']);
+        expect(Scheduler).toFlushAndYield(['only:a-one', 'Sync effect']);
         ReactNoop.flushPassiveEffects();
         expect(sourceA.listenerCount).toBe(1);
 
@@ -630,7 +630,7 @@ describe('useMutableSource', () => {
     });
 
     it('should only update components whose subscriptions fire', () => {
-      const source = createComplexSource('one', 'one');
+      const source = createComplexSource('a:one', 'b:one');
       const mutableSource = createMutableSource(source);
 
       // Subscribe to part of the store.
@@ -672,7 +672,7 @@ describe('useMutableSource', () => {
     });
 
     it('should detect tearing in part of the store not yet subscribed to', () => {
-      const source = createComplexSource('one', 'one');
+      const source = createComplexSource('a:one', 'b:one');
       const mutableSource = createMutableSource(source);
 
       // Subscribe to part of the store.
@@ -1054,32 +1054,31 @@ describe('useMutableSource', () => {
         };
       }
 
-      function App({toggle}) {
+      function App({getSnapshot}) {
         const state = useMutableSource(
           mutableSource,
-          toggle ? getSnapshotB : getSnapshotA,
+          getSnapshot,
           defaultSubscribe,
         );
-        const result = (toggle ? 'on: ' : 'off: ') + state;
-        return result;
+        return state;
       }
 
       const root = ReactNoop.createRoot();
       await act(async () => {
-        root.render(<App toggle={false} />);
+        root.render(<App getSnapshot={getSnapshotA} />);
       });
-      expect(root).toMatchRenderedOutput('off: initial');
+      expect(root).toMatchRenderedOutput('initial');
 
       await act(async () => {
         mutateB('Updated B');
-        root.render(<App toggle={true} />);
+        root.render(<App getSnapshot={getSnapshotB} />);
       });
-      expect(root).toMatchRenderedOutput('on: Updated B');
+      expect(root).toMatchRenderedOutput('Updated B');
 
       await act(async () => {
         mutateB('Another update');
       });
-      expect(root).toMatchRenderedOutput('on: Another update');
+      expect(root).toMatchRenderedOutput('Another update');
     });
 
     it('should clear the update queue when getSnapshot changes with pending lower priority updates', async () => {
@@ -1238,7 +1237,7 @@ describe('useMutableSource', () => {
         ReactNoop.discreteUpdates(() => {
           // At high priority, toggle y so that it reads from A instead of B.
           // Simultaneously, mutate A.
-          mutateA('high-pri baz');
+          mutateA('baz');
           root.render(
             <App
               getSnapshotFirst={getSnapshotA}
@@ -1247,7 +1246,7 @@ describe('useMutableSource', () => {
           );
 
           // If this update were processed before the next mutation,
-          // it would be expected to yield "high-pri baz" and "high-pri baz".
+          // it would be expected to yield "baz" and "baz".
         });
 
         // At lower priority, mutate A again.
@@ -1265,6 +1264,316 @@ describe('useMutableSource', () => {
       expect(Scheduler).toHaveYielded(['x: bar, y: bar']);
     });
 
+    it('getSnapshot changes and then source is mutated in between paint and passive effect phase', async () => {
+      const source = createSource({
+        a: 'foo',
+        b: 'bar',
+      });
+      const mutableSource = createMutableSource(source);
+
+      function mutateB(newB) {
+        source.value = {
+          ...source.value,
+          b: newB,
+        };
+      }
+
+      const getSnapshotA = () => source.value.a;
+      const getSnapshotB = () => source.value.b;
+
+      function App({getSnapshot}) {
+        const value = useMutableSource(
+          mutableSource,
+          getSnapshot,
+          defaultSubscribe,
+        );
+
+        Scheduler.unstable_yieldValue('Render: ' + value);
+        React.useEffect(() => {
+          Scheduler.unstable_yieldValue('Commit: ' + value);
+        }, [value]);
+
+        return value;
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(async () => {
+        root.render(<App getSnapshot={getSnapshotA} />);
+      });
+      expect(Scheduler).toHaveYielded(['Render: foo', 'Commit: foo']);
+
+      await act(async () => {
+        // Switch getSnapshot to read from B instead
+        root.render(<App getSnapshot={getSnapshotB} />);
+        // Render and finish the tree, but yield right after paint, before
+        // the passive effects have fired.
+        expect(Scheduler).toFlushUntilNextPaint(['Render: bar']);
+        // Then mutate B.
+        mutateB('baz');
+      });
+      expect(Scheduler).toHaveYielded([
+        // Fires the effect from the previous render
+        'Commit: bar',
+        // During that effect, it should detect that the snapshot has changed
+        // and re-render.
+        'Render: baz',
+        'Commit: baz',
+      ]);
+      expect(root).toMatchRenderedOutput('baz');
+    });
+
+    it('getSnapshot changes and then source is mutated in between paint and passive effect phase, case 2', async () => {
+      const source = createSource({
+        a: 'a0',
+        b: 'b0',
+      });
+      const mutableSource = createMutableSource(source);
+
+      const getSnapshotA = () => source.value.a;
+      const getSnapshotB = () => source.value.b;
+
+      function mutateA(newA) {
+        source.value = {
+          ...source.value,
+          a: newA,
+        };
+      }
+
+      function App({getSnapshotFirst, getSnapshotSecond}) {
+        const first = useMutableSource(
+          mutableSource,
+          getSnapshotFirst,
+          defaultSubscribe,
+        );
+        const second = useMutableSource(
+          mutableSource,
+          getSnapshotSecond,
+          defaultSubscribe,
+        );
+
+        return `first: ${first}, second: ${second}`;
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(async () => {
+        root.render(
+          <App
+            getSnapshotFirst={getSnapshotA}
+            getSnapshotSecond={getSnapshotB}
+          />,
+        );
+      });
+      expect(root.getChildrenAsJSX()).toEqual('first: a0, second: b0');
+
+      await act(async () => {
+        // Switch the second getSnapshot to also read from A
+        root.render(
+          <App
+            getSnapshotFirst={getSnapshotA}
+            getSnapshotSecond={getSnapshotA}
+          />,
+        );
+        // Render and finish the tree, but yield right after paint, before
+        // the passive effects have fired.
+        expect(Scheduler).toFlushUntilNextPaint([]);
+
+        // Now mutate A. Both hooks should update.
+        // This is at high priority so that it doesn't get batched with default
+        // priority updates that might fire during the passive effect
+        ReactNoop.discreteUpdates(() => {
+          mutateA('a1');
+        });
+        expect(Scheduler).toFlushUntilNextPaint([]);
+
+        expect(root.getChildrenAsJSX()).toEqual('first: a1, second: a1');
+      });
+
+      expect(root.getChildrenAsJSX()).toEqual('first: a1, second: a1');
+    });
+
+    it('getSnapshot changes and then source is mutated during interleaved event', async () => {
+      const {useEffect} = React;
+
+      const source = createComplexSource('1', '2');
+      const mutableSource = createMutableSource(source);
+
+      // Subscribe to part of the store.
+      const getSnapshotA = s => s.valueA;
+      const subscribeA = (s, callback) => s.subscribeA(callback);
+      const configA = [getSnapshotA, subscribeA];
+
+      const getSnapshotB = s => s.valueB;
+      const subscribeB = (s, callback) => s.subscribeB(callback);
+      const configB = [getSnapshotB, subscribeB];
+
+      function App({parentConfig, childConfig}) {
+        const [getSnapshot, subscribe] = parentConfig;
+        const parentValue = useMutableSource(
+          mutableSource,
+          getSnapshot,
+          subscribe,
+        );
+
+        Scheduler.unstable_yieldValue('Parent: ' + parentValue);
+
+        return (
+          <Child
+            parentConfig={parentConfig}
+            childConfig={childConfig}
+            parentValue={parentValue}
+          />
+        );
+      }
+
+      function Child({parentConfig, childConfig, parentValue}) {
+        const [getSnapshot, subscribe] = childConfig;
+        const childValue = useMutableSource(
+          mutableSource,
+          getSnapshot,
+          subscribe,
+        );
+
+        Scheduler.unstable_yieldValue('Child: ' + childValue);
+
+        let result = `${parentValue}, ${childValue}`;
+
+        if (parentConfig === childConfig) {
+          // When both components read using the same config, the two values
+          // must be consistent.
+          if (parentValue !== childValue) {
+            result = 'Oops, tearing!';
+          }
+        }
+
+        useEffect(() => {
+          Scheduler.unstable_yieldValue('Commit: ' + result);
+        }, [result]);
+
+        return result;
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(async () => {
+        root.render(<App parentConfig={configA} childConfig={configB} />);
+      });
+      expect(Scheduler).toHaveYielded([
+        'Parent: 1',
+        'Child: 2',
+        'Commit: 1, 2',
+      ]);
+
+      await act(async () => {
+        // Switch the parent and the child to read using the same config
+        root.render(<App parentConfig={configB} childConfig={configB} />);
+        // Start rendering the parent, but yield before rendering the child
+        expect(Scheduler).toFlushAndYieldThrough(['Parent: 2']);
+
+        // Mutate the config. This is at lower priority so that 1) to make sure
+        // it doesn't happen to get batched with the in-progress render, and 2)
+        // so it doesn't interrupt the in-progress render.
+        Scheduler.unstable_runWithPriority(
+          Scheduler.unstable_IdlePriority,
+          () => {
+            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',
+
+        // Eventually the child corrects itself, because of the check that
+        // occurs when re-subscribing.
+        'Child: 3',
+        'Commit: 3, 3',
+      ]);
+    });
+
+    it('should not tear with newly mounted component when updates were scheduled at a lower priority', async () => {
+      const source = createSource('one');
+      const mutableSource = createMutableSource(source);
+
+      let committedA = null;
+      let committedB = null;
+
+      const onRender = () => {
+        if (committedB !== null) {
+          expect(committedA).toBe(committedB);
+        }
+      };
+
+      function ComponentA() {
+        const snapshot = useMutableSource(
+          mutableSource,
+          defaultGetSnapshot,
+          defaultSubscribe,
+        );
+        Scheduler.unstable_yieldValue(`a:${snapshot}`);
+        React.useEffect(() => {
+          committedA = snapshot;
+        }, [snapshot]);
+        return <div>{`a:${snapshot}`}</div>;
+      }
+      function ComponentB() {
+        const snapshot = useMutableSource(
+          mutableSource,
+          defaultGetSnapshot,
+          defaultSubscribe,
+        );
+        Scheduler.unstable_yieldValue(`b:${snapshot}`);
+        React.useEffect(() => {
+          committedB = snapshot;
+        }, [snapshot]);
+        return <div>{`b:${snapshot}`}</div>;
+      }
+
+      // Mount ComponentA with data version 1
+      act(() => {
+        ReactNoop.render(
+          <React.Profiler id="root" onRender={onRender}>
+            <ComponentA />
+          </React.Profiler>,
+          () => Scheduler.unstable_yieldValue('Sync effect'),
+        );
+      });
+      expect(Scheduler).toHaveYielded(['a:one', 'Sync effect']);
+      expect(source.listenerCount).toBe(1);
+
+      // Mount ComponentB with version 1 (but don't commit it)
+      act(() => {
+        ReactNoop.render(
+          <React.Profiler id="root" onRender={onRender}>
+            <ComponentA />
+            <ComponentB />
+          </React.Profiler>,
+          () => Scheduler.unstable_yieldValue('Sync effect'),
+        );
+        expect(Scheduler).toFlushAndYieldThrough([
+          'a:one',
+          'b:one',
+          'Sync effect',
+        ]);
+        expect(source.listenerCount).toBe(1);
+
+        // Mutate -> schedule update for ComponentA
+        Scheduler.unstable_runWithPriority(
+          Scheduler.unstable_IdlePriority,
+          () => {
+            source.value = 'two';
+          },
+        );
+
+        // Commit ComponentB -> notice the change and schedule an update for ComponentB
+        expect(Scheduler).toFlushAndYield(['a:two', 'b:two']);
+        expect(source.listenerCount).toBe(2);
+      });
+    });
+
     if (__DEV__) {
       describe('dev warnings', () => {
         it('should warn if the subscribe function does not return an unsubscribe function', () => {