Added additional useMutableSource test and bugfix

packages/react-reconciler/src/ReactFiberHooks.js

@@ -885,6 +885,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>,
@@ -998,18 +999,36 @@ function useMutableSource<Source, Snapshot>(
 
   // Sync the values needed by our subscribe function after each commit.
   dispatcher.useEffect(() => {
+    const didGetSnapshotChange = !is(refs.getSnapshot, getSnapshot);
     refs.getSnapshot = getSnapshot;
 
-    // Because getSnapshot is shared with subscriptions via a ref,
-    // we don't resubscribe when getSnapshot changes.
-    // This means that we also don't check for any missed mutations
-    // between the render and the passive commit though.
-    // So we need to check here, just like when we newly subscribe.
-    const maybeNewVersion = getVersion(source._source);
-    if (!is(version, maybeNewVersion)) {
-      const maybeNewSnapshot = getSnapshot(source._source);
-      if (!is(snapshot, maybeNewSnapshot)) {
-        setSnapshot(maybeNewSnapshot);
+    // Normally the dispatch function for a state hook never changes,
+    // but in the case of this hook, it will change if getSnapshot changes.
+    // In that case, the subscription below will have cloesd over the previous function,
+    // so we use a ref to ensure that handleChange() always has the latest version.
+    refs.setSnapshot = setSnapshot;
+
+    // This effect runs on mount, even though getSnapshot hasn't changed.
+    // In that case we can avoid the additional checks for a changed snapshot,
+    // because the subscription effect below will cover this.
+    if (didGetSnapshotChange) {
+      // Because getSnapshot is shared with subscriptions via a ref,
+      // we don't resubscribe when getSnapshot changes.
+      // This means that we also don't check for any missed mutations
+      // between the render and the passive commit though.
+      // So we need to check here, just like when we newly subscribe.
+      const maybeNewVersion = getVersion(source._source);
+      if (!is(version, maybeNewVersion)) {
+        const maybeNewSnapshot = getSnapshot(source._source);
+        if (!is(snapshot, maybeNewSnapshot)) {
+          // If the source mutated between render and now,
+          // we may have a state update already scheduled from the old getSnapshot.
+          // It's important to ensure that this update is included with that one.
+          // TODO This sucks; what's the right way of doing this?
+          runWithPriority(UserBlockingPriority, () => {
+            setSnapshot(maybeNewSnapshot);
+          });
+        }
       }
     }
   }, [getSnapshot]);
@@ -1020,8 +1039,10 @@ function useMutableSource<Source, Snapshot>(
   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();
@@ -1038,9 +1059,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),
+        );
       }
     };
 
@@ -1076,11 +1099,31 @@ 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.
+  const didGetSnapshotChange = !is(prevGetSnapshot, getSnapshot);
   if (
     !is(prevSource, source) ||
     !is(prevSubscribe, subscribe) ||
-    !is(prevGetSnapshot, getSnapshot)
+    didGetSnapshotChange
   ) {
+    if (didGetSnapshotChange) {
+      // Create a new queue and setState method,
+      // So if there are interleaved updates, they get pushed to the older queue.
+      const newQueue = {
+        pending: null,
+        dispatch: null,
+        lastRenderedReducer: basicStateReducer,
+        lastRenderedState: snapshot,
+      };
+      newQueue.dispatch = setSnapshot = (dispatchAction.bind(
+        null,
+        currentlyRenderingFiber,
+        newQueue,
+      ): any);
+      stateHook.queue = newQueue;
+    }
+
+    // Now when this becomes current, the previous queue and dispatch method
+    // are complete discarded, including any interleaving updates that occur.
     stateHook.baseQueue = null;
     snapshot = readFromUnsubcribedMutableSource(root, source, getSnapshot);
     stateHook.memoizedState = stateHook.baseState = snapshot;
@@ -1098,6 +1141,7 @@ function mountMutableSource<Source, Snapshot>(
   hook.memoizedState = ({
     refs: {
       getSnapshot,
+      setSnapshot: (null: any),
     },
     source,
     subscribe,

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

@@ -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 () => {
@@ -1323,6 +1322,86 @@ describe('useMutableSource', () => {
       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 = getSnapshotFirst(source);
+        const first = useMutableSource(
+          mutableSource,
+          getSnapshotFirst,
+          defaultSubscribe,
+        );
+        const second = useMutableSource(
+          mutableSource,
+          getSnapshotSecond,
+          defaultSubscribe,
+        );
+
+        let result = `first: ${first}, second: ${second}`;
+
+        if (getSnapshotFirst === getSnapshotSecond) {
+          // When both getSnapshot functions are equal,
+          // the two values must be consistent.
+          if (first !== second) {
+            //result = 'Oops, tearing!';
+          }
+        }
+
+        return result;
+      }
+
+      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()).not.toEqual('Oops, tearing!');
+        expect(root.getChildrenAsJSX()).toEqual('first: a1, second: a1');
+      });
+
+      expect(root.getChildrenAsJSX()).toEqual('first: a1, second: a1');
+    });
+
     if (__DEV__) {
       describe('dev warnings', () => {
         it('should warn if the subscribe function does not return an unsubscribe function', () => {