useMutableSource: bugfix for new getSnapshot with mutation

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;
@@ -1000,7 +1001,49 @@ 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;
+
+    // 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)) {
+          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]);
 
   // If we got a new source or subscribe function,
@@ -1009,8 +1052,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();
@@ -1027,9 +1072,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),
+        );
       }
     };
 
@@ -1048,6 +1095,20 @@ function useMutableSource<Source, Snapshot>(
       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);
+
+        // We missed a mutation before committing.
+        // It's possible that other components using this source also have pending updates scheduled.
+        // In that case, we should ensure they all commit together.
+        markRootExpiredAtTime(root, getLastPendingExpirationTime(root));
       }
     }
 
@@ -1065,11 +1126,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.
+      // 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 +1168,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 {