Fix useMutableSource tearing bug

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

@@ -997,14 +997,11 @@ function useMutableSource<Source, Snapshot>(
         const suspenseConfig = requestCurrentSuspenseConfig();
         const lane = requestUpdateLane(fiber, suspenseConfig);
         markRootMutableRead(root, lane);
-
-        // If the source mutated between render and now,
-        // 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);
       }
+      // If the source mutated between render and now,
+      // there may be state updates already scheduled from the old source.
+      // Entangle the updates so that they render in the same batch.
+      markRootEntangled(root, root.mutableReadLanes);
     }
   }, [getSnapshot, source, subscribe]);
 

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

@@ -986,15 +986,14 @@ function useMutableSource<Source, Snapshot>(
           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));
       }
+      // 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]);
 

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

@@ -1398,6 +1398,101 @@ describe('useMutableSource', () => {
     expect(root.getChildrenAsJSX()).toEqual('first: a1, second: a1');
   });
 
+  // @gate experimental
+  it(
+    'if source is mutated after initial read but before subscription is set ' +
+      'up, should still entangle all pending mutations even if snapshot of ' +
+      'new subscription happens to match',
+    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 mutateB(newB) {
+        source.value = {
+          ...source.value,
+          b: newB,
+        };
+      }
+
+      function Read({getSnapshot}) {
+        const value = useMutableSource(
+          mutableSource,
+          getSnapshot,
+          defaultSubscribe,
+        );
+        Scheduler.unstable_yieldValue(value);
+        return value;
+      }
+
+      function Text({text}) {
+        Scheduler.unstable_yieldValue(text);
+        return text;
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(async () => {
+        root.render(
+          <>
+            <Read getSnapshot={getSnapshotA} />
+          </>,
+        );
+      });
+      expect(Scheduler).toHaveYielded(['a0']);
+      expect(root).toMatchRenderedOutput('a0');
+
+      await act(async () => {
+        root.render(
+          <>
+            <Read getSnapshot={getSnapshotA} />
+            <Read getSnapshot={getSnapshotB} />
+            <Text text="c" />
+          </>,
+        );
+
+        expect(Scheduler).toFlushAndYieldThrough(['a0', 'b0']);
+        // Mutate in an event. This schedules a subscription update on a, which
+        // already mounted, but not b, which hasn't subscribed yet.
+        mutateA('a1');
+        mutateB('b1');
+
+        // Mutate again at lower priority. This will schedule another subscription
+        // update on a, but not b. When b mounts and subscriptions, the value it
+        // read during render will happen to match the latest value. But it should
+        // still entangle the updates to prevent the previous update (a1) from
+        // rendering by itself.
+        Scheduler.unstable_runWithPriority(
+          Scheduler.unstable_IdlePriority,
+          () => {
+            mutateA('a0');
+            mutateB('b0');
+          },
+        );
+        // Finish the current render
+        expect(Scheduler).toFlushUntilNextPaint(['c']);
+        // a0 will re-render because of the mutation update. But it should show
+        // the latest value, not the intermediate one, to avoid tearing with b.
+        expect(Scheduler).toFlushUntilNextPaint(['a0']);
+        expect(root).toMatchRenderedOutput('a0b0c');
+        // We should be done.
+        expect(Scheduler).toFlushAndYield([]);
+        expect(root).toMatchRenderedOutput('a0b0c');
+      });
+    },
+  );
+
   // @gate experimental
   it('getSnapshot changes and then source is mutated during interleaved event', async () => {
     const {useEffect} = React;