useMutableSource reduce amount of metadata cached during render

Author: Brian Vaughn

packages/react-reconciler/src/ReactFiberHooks.js

@@ -23,6 +23,7 @@ import type {
 } from 'shared/ReactMutableSource';
 
 import ReactSharedInternals from 'shared/ReactSharedInternals';
+import {HostRoot} from 'shared/ReactWorkTags';
 
 import {NoWork, Sync} from './ReactFiberExpirationTime';
 import {readContext} from './ReactFiberNewContext';
@@ -45,7 +46,8 @@ import {
   warnIfNotScopedWithMatchingAct,
   markRenderEventTimeAndConfig,
   markUnprocessedUpdateTime,
-  getMutableSourceMetadata,
+  getMutableSourcePendingExpirationTime,
+  warnAboutUpdateOnUnmountedFiberInDEV,
 } from './ReactFiberWorkLoop';
 
 import invariant from 'shared/invariant';
@@ -897,13 +899,15 @@ function useMutableSourceImpl<S>(
     }
   }
 
-  const metadata = getMutableSourceMetadata(source);
+  const pendingExpirationTime = getMutableSourcePendingExpirationTime(source);
 
-  // Is it safe to read from this source during the current render?
-  // If the source has not yet been subscribed to, we can use the version number to determine this.
-  // Else we can use the expiration time as an indicator of any future scheduled updates.
   let isSafeToReadFromSource = false;
-  if (metadata.subscriptionCount === 0) {
+
+  // Is it safe to read from this source during the current render?
+  // 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.
+  // If there are no pending updates, we can use the work-in-progress version.
+  if (pendingExpirationTime === null) {
     const lastReadVersion = getWorkInProgressVersion(source);
     if (lastReadVersion === null) {
       // This is the only case where we need to actually update the version number.
@@ -926,8 +930,8 @@ function useMutableSourceImpl<S>(
     );
 
     isSafeToReadFromSource =
-      metadata.expirationTime === NoWork ||
-      metadata.expirationTime >= expirationTime;
+      pendingExpirationTime === NoWork ||
+      pendingExpirationTime >= expirationTime;
   }
 
   let prevMemoizedState = ((hook.memoizedState: any): ?MutableSourceState<S>);
@@ -977,22 +981,43 @@ function useMutableSourceImpl<S>(
 
     const create = () => {
       const scheduleUpdate = () => {
-        const currentTime = requestCurrentTimeForUpdate();
-        const suspenseConfig = requestCurrentSuspenseConfig();
-        const expirationTime = computeExpirationForFiber(
-          currentTime,
-          fiber,
-          suspenseConfig,
-        );
+        let node = fiber;
+        let root = null;
+        while (node !== null) {
+          if (node.tag === HostRoot) {
+            root = node.stateNode;
+            break;
+          }
+          node = node.return;
+        }
 
-        // Make sure reads during future renders will know there's a pending update.
-        // This will prevent a higher priority update from reading a newer version of the source,
-        // and causing a tear between that render and previous renders.
-        if (expirationTime > metadata.expirationTime) {
-          metadata.expirationTime = expirationTime;
+        if (root === null) {
+          warnAboutUpdateOnUnmountedFiberInDEV(fiber);
+          return;
         }
 
-        scheduleWork(fiber, expirationTime);
+        const alreadyScheduledExpirationTime = root.mutableSourcePendingUpdateMap.get(
+          source,
+        );
+
+        // If an update is already scheduled for this source, re-use the same priority.
+        if (alreadyScheduledExpirationTime !== undefined) {
+          scheduleWork(fiber, alreadyScheduledExpirationTime);
+        } else {
+          const currentTime = requestCurrentTimeForUpdate();
+          const suspenseConfig = requestCurrentSuspenseConfig();
+          const expirationTime = computeExpirationForFiber(
+            currentTime,
+            fiber,
+            suspenseConfig,
+          );
+          scheduleWork(fiber, expirationTime);
+
+          // Make sure reads during future renders will know there's a pending update.
+          // This will prevent a higher priority update from reading a newer version of the source,
+          // and causing a tear between that render and previous renders.
+          root.mutableSourcePendingUpdateMap.set(source, expirationTime);
+        }
       };
 
       // Was the source mutated between when we rendered and when we're subscribing?
@@ -1002,16 +1027,8 @@ function useMutableSourceImpl<S>(
         scheduleUpdate();
       }
 
-      const unsubscribe = subscribe(scheduleUpdate);
-      metadata.subscriptionCount++;
-
-      memoizedState.destroy = () => {
-        metadata.subscriptionCount--;
-
-        // TODO (useMutableSource) If count is 0, flag this source for possible cleanup.
-
-        unsubscribe();
-      };
+      // Unsubscribe on destroy.
+      memoizedState.destroy = subscribe(scheduleUpdate);
 
       return memoizedState.destroy;
     };

packages/react-reconciler/src/ReactFiberRoot.js

@@ -15,7 +15,7 @@ import type {Thenable} from './ReactFiberWorkLoop';
 import type {Interaction} from 'scheduler/src/Tracing';
 import type {SuspenseHydrationCallbacks} from './ReactFiberSuspenseComponent';
 import type {ReactPriorityLevel} from './SchedulerWithReactIntegration';
-import type {MutableSourceMetadataMap} from 'shared/ReactMutableSource';
+import type {MutableSourcePendingUpdateMap} from 'shared/ReactMutableSource';
 
 import {noTimeout} from './ReactFiberHostConfig';
 import {createHostRootFiber} from './ReactFiber';
@@ -77,7 +77,7 @@ type BaseFiberRootProperties = {|
   lastExpiredTime: ExpirationTime,
   // Used by useMutableSource hook to avoid tearing within this root
   // when external, mutable sources are read from during render.
-  mutableSourceMetadata: MutableSourceMetadataMap,
+  mutableSourcePendingUpdateMap: MutableSourcePendingUpdateMap,
 |};
 
 // The following attributes are only used by interaction tracing builds.
@@ -127,7 +127,7 @@ function FiberRootNode(containerInfo, tag, hydrate) {
   this.nextKnownPendingLevel = NoWork;
   this.lastPingedTime = NoWork;
   this.lastExpiredTime = NoWork;
-  this.mutableSourceMetadata = new Map();
+  this.mutableSourcePendingUpdateMap = new Map();
 
   if (enableSchedulerTracing) {
     this.interactionThreadID = unstable_getThreadID();

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -14,10 +14,7 @@ import type {ReactPriorityLevel} from './SchedulerWithReactIntegration';
 import type {Interaction} from 'scheduler/src/Tracing';
 import type {SuspenseConfig} from './ReactFiberSuspenseConfig';
 import type {SuspenseState} from './ReactFiberSuspenseComponent';
-import type {
-  MutableSource,
-  MutableSourceMetadata,
-} from 'shared/ReactMutableSource';
+import type {MutableSource} from 'shared/ReactMutableSource';
 
 import {
   initializeWorkInProgressVersionMap as initializeMutableSourceWorkInProgressVersionMap,
@@ -299,24 +296,14 @@ let spawnedWorkDuringRender: null | Array<ExpirationTime> = null;
 // receive the same expiration time. Otherwise we get tearing.
 let currentEventTime: ExpirationTime = NoWork;
 
-export function getMutableSourceMetadata(
+export function getMutableSourcePendingExpirationTime(
   source: MutableSource,
-): MutableSourceMetadata {
+): ExpirationTime | null {
   invariant(workInProgressRoot !== null, 'Expected a work-in-progress root.');
-
-  let metadata = workInProgressRoot.mutableSourceMetadata.get(source);
-  if (metadata !== undefined) {
-    return metadata;
-  } else {
-    metadata = {
-      expirationTime: NoWork,
-      subscriptionCount: 0,
-    };
-
-    workInProgressRoot.mutableSourceMetadata.set(source, metadata);
-
-    return ((metadata: any): MutableSourceMetadata);
-  }
+  const expirationTime = workInProgressRoot.mutableSourcePendingUpdateMap.get(
+    source,
+  );
+  return expirationTime !== undefined ? expirationTime : null;
 }
 
 export function requestCurrentTimeForUpdate() {
@@ -402,10 +389,7 @@ export function computeExpirationForFiber(
   return expirationTime;
 }
 
-export function scheduleUpdateOnFiber(
-  fiber: Fiber,
-  expirationTime: ExpirationTime,
-) {
+function scheduleUpdateOnFiber(fiber: Fiber, expirationTime: ExpirationTime) {
   checkForNestedUpdates();
   warnAboutInvalidUpdatesOnClassComponentsInDEV(fiber);
 
@@ -2027,6 +2011,15 @@ function commitRootImpl(root, renderPriorityLevel) {
     nestedUpdateCount = 0;
   }
 
+  // Remove pending mutable source entries that we've completed processing.
+  root.mutableSourcePendingUpdateMap.forEach(
+    (pendingExpirationTime, source) => {
+      if (pendingExpirationTime <= expirationTime) {
+        root.mutableSourcePendingUpdateMap.delete(source);
+      }
+    },
+  );
+
   resetMutableSourceWorkInProgressVersionMap();
 
   onCommitRoot(finishedWork.stateNode, expirationTime);
@@ -2280,9 +2273,6 @@ function flushPassiveEffectsImpl() {
     finishPendingInteractions(root, expirationTime);
   }
 
-  // TODO (useMutableSource) Remove metadata for mutable sources that are no longer in use.
-  // This check comes after passive effects, because that's when sources are unsubscribed from.
-
   executionContext = prevExecutionContext;
 
   flushSyncCallbackQueue();
@@ -2626,7 +2616,7 @@ function checkForInterruption(
 }
 
 let didWarnStateUpdateForUnmountedComponent: Set<string> | null = null;
-function warnAboutUpdateOnUnmountedFiberInDEV(fiber) {
+export function warnAboutUpdateOnUnmountedFiberInDEV(fiber: Fiber) {
   if (__DEV__) {
     const tag = fiber.tag;
     if (

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

@@ -2098,8 +2098,6 @@ describe('ReactHooksWithNoopRenderer', () => {
       });
       expect(Scheduler).toHaveYielded(['a:two', 'b:two', 'Sync effect']);
     });
-
-    // TODO (useMutableSource) Edge case: make sure we don't leak on root Map (how to test this without internals?)
   });
 
   describe('useCallback', () => {

packages/shared/ReactMutableSource.js

@@ -20,23 +20,12 @@ export type MutableSourceHookConfig<S> = {|
   subscribe: (callback: Function) => () => void,
 |};
 
-export type MutableSourceMetadata = {|
-  // Expiration time of most recently scheduled update.
-  // Used to determine if a source is safe to read during updates.
-  // If the render’s expiration time is ≤ this value,
-  // the source has not changed since the last render and is safe to read from.
-  expirationTime: ExpirationTime,
-
-  // Number of hooks that are subscribed as of the most recently committed render.
-  // This value is used to determine when a source is no longer in use,
-  // and should be removed from the root map to avoid a memory leak.
-  subscriptionCount: number,
-|};
-
-export type MutableSourceMetadataMap = Map<
-  MutableSource,
-  MutableSourceMetadata,
->;
+// Tracks expiration time for all mutable sources with pending updates.
+// Used to determine if a source is safe to read during updates.
+// If there are no entries in this map for a given source,
+// or if the current render’s expiration time is ≤ this value,
+// it is safe to read from the source without tearing.
+export type MutableSourcePendingUpdateMap = Map<MutableSource, ExpirationTime>;
 
 // Tracks the version of each source at the time it was most recently read.
 // Used to determine if a source is safe to read from before it has been subscribed to.