Memoize promise listeners to prevent exponential growth (#14429)

* Memoize promise listeners to prevent exponential growth

Previously, React would attach a new listener every time a promise is
thrown, regardless of whether the same listener was already attached
during a previous render. Because React attempts to render every time
a promise resolves, the number of listeners grows quickly.

This was especially bad in synchronous mode because the renders that
happen when the promise pings are not batched together. So if a single
promise has multiple listeners for the same root, there will be multiple
renders, which in turn results in more listeners being added to the
remaining unresolved promises. This results in exponential growth in
the number of listeners with respect to the number of IO-bound
components in a single render.

Fixes #14220

* Memoize on the root and Suspense fiber instead of on the promise

* Add TODO to fix persistent mode tests

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -1474,6 +1474,7 @@ function updateSuspenseComponent(
         );
       }
     }
+    workInProgress.stateNode = current.stateNode;
   }
 
   workInProgress.memoizedState = nextState;

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -20,7 +20,9 @@ import type {ExpirationTime} from './ReactFiberExpirationTime';
 import type {CapturedValue, CapturedError} from './ReactCapturedValue';
 import type {SuspenseState} from './ReactFiberSuspenseComponent';
 import type {FunctionComponentUpdateQueue} from './ReactFiberHooks';
+import type {Thenable} from './ReactFiberScheduler';
 
+import {unstable_wrap as Schedule_tracing_wrap} from 'scheduler/tracing';
 import {
   enableHooks,
   enableSchedulerTracing,
@@ -88,6 +90,7 @@ import {
 import {
   captureCommitPhaseError,
   requestCurrentTime,
+  retryTimedOutBoundary,
 } from './ReactFiberScheduler';
 import {
   NoEffect as NoHookEffect,
@@ -106,6 +109,8 @@ if (__DEV__) {
   didWarnAboutUndefinedSnapshotBeforeUpdate = new Set();
 }
 
+const PossiblyWeakSet = typeof WeakSet === 'function' ? WeakSet : Set;
+
 export function logError(boundary: Fiber, errorInfo: CapturedValue<mixed>) {
   const source = errorInfo.source;
   let stack = errorInfo.stack;
@@ -1180,6 +1185,30 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
       if (primaryChildParent !== null) {
         hideOrUnhideAllChildren(primaryChildParent, newDidTimeout);
       }
+
+      // If this boundary just timed out, then it will have a set of thenables.
+      // For each thenable, attach a listener so that when it resolves, React
+      // attempts to re-render the boundary in the primary (pre-timeout) state.
+      const thenables: Set<Thenable> | null = (finishedWork.updateQueue: any);
+      if (thenables !== null) {
+        finishedWork.updateQueue = null;
+        let retryCache = finishedWork.stateNode;
+        if (retryCache === null) {
+          retryCache = finishedWork.stateNode = new PossiblyWeakSet();
+        }
+        thenables.forEach(thenable => {
+          // Memoize using the boundary fiber to prevent redundant listeners.
+          let retry = retryTimedOutBoundary.bind(null, finishedWork, thenable);
+          if (enableSchedulerTracing) {
+            retry = Schedule_tracing_wrap(retry);
+          }
+          if (!retryCache.has(thenable)) {
+            retryCache.add(thenable);
+            thenable.then(retry, retry);
+          }
+        });
+      }
+
       return;
     }
     case IncompleteClassComponent: {

packages/react-reconciler/src/ReactFiberPendingPriority.js

@@ -62,6 +62,10 @@ export function markCommittedPriorityLevels(
     return;
   }
 
+  if (earliestRemainingTime < root.latestPingedTime) {
+    root.latestPingedTime = NoWork;
+  }
+
   // Let's see if the previous latest known pending level was just flushed.
   const latestPendingTime = root.latestPendingTime;
   if (latestPendingTime !== NoWork) {
@@ -209,10 +213,8 @@ export function markPingedPriorityLevel(
 }
 
 function clearPing(root, completedTime) {
-  // TODO: Track whether the root was pinged during the render phase. If so,
-  // we need to make sure we don't lose track of it.
   const latestPingedTime = root.latestPingedTime;
-  if (latestPingedTime !== NoWork && latestPingedTime >= completedTime) {
+  if (latestPingedTime >= completedTime) {
     root.latestPingedTime = NoWork;
   }
 }

packages/react-reconciler/src/ReactFiberRoot.js

@@ -10,6 +10,7 @@
 import type {Fiber} from './ReactFiber';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 import type {TimeoutHandle, NoTimeout} from './ReactFiberHostConfig';
+import type {Thenable} from './ReactFiberScheduler';
 import type {Interaction} from 'scheduler/src/Tracing';
 
 import {noTimeout} from './ReactFiberHostConfig';
@@ -51,6 +52,11 @@ type BaseFiberRootProperties = {|
   // be retried.
   latestPingedTime: ExpirationTime,
 
+  pingCache:
+    | WeakMap<Thenable, Set<ExpirationTime>>
+    | Map<Thenable, Set<ExpirationTime>>
+    | null,
+
   // If an error is thrown, and there are no more updates in the queue, we try
   // rendering from the root one more time, synchronously, before handling
   // the error.
@@ -121,6 +127,8 @@ export function createFiberRoot(
       latestSuspendedTime: NoWork,
       latestPingedTime: NoWork,
 
+      pingCache: null,
+
       didError: false,
 
       pendingCommitExpirationTime: NoWork,
@@ -144,6 +152,8 @@ export function createFiberRoot(
       containerInfo: containerInfo,
       pendingChildren: null,
 
+      pingCache: null,
+
       earliestPendingTime: NoWork,
       latestPendingTime: NoWork,
       earliestSuspendedTime: NoWork,

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -125,13 +125,8 @@ import {
   computeAsyncExpiration,
   computeInteractiveExpiration,
 } from './ReactFiberExpirationTime';
-import {ConcurrentMode, ProfileMode, NoContext} from './ReactTypeOfMode';
-import {
-  enqueueUpdate,
-  resetCurrentlyProcessingQueue,
-  ForceUpdate,
-  createUpdate,
-} from './ReactUpdateQueue';
+import {ConcurrentMode, ProfileMode} from './ReactTypeOfMode';
+import {enqueueUpdate, resetCurrentlyProcessingQueue} from './ReactUpdateQueue';
 import {createCapturedValue} from './ReactCapturedValue';
 import {
   isContextProvider as isLegacyContextProvider,
@@ -1646,62 +1641,62 @@ function renderDidError() {
   nextRenderDidError = true;
 }
 
-function retrySuspendedRoot(
+function pingSuspendedRoot(
   root: FiberRoot,
-  boundaryFiber: Fiber,
-  sourceFiber: Fiber,
-  suspendedTime: ExpirationTime,
+  thenable: Thenable,
+  pingTime: ExpirationTime,
 ) {
-  let retryTime;
+  // A promise that previously suspended React from committing has resolved.
+  // If React is still suspended, try again at the previous level (pingTime).
 
-  if (isPriorityLevelSuspended(root, suspendedTime)) {
-    // Ping at the original level
-    retryTime = suspendedTime;
+  const pingCache = root.pingCache;
+  if (pingCache !== null) {
+    // The thenable resolved, so we no longer need to memoize, because it will
+    // never be thrown again.
+    pingCache.delete(thenable);
+  }
 
-    markPingedPriorityLevel(root, retryTime);
+  if (nextRoot !== null && nextRenderExpirationTime === pingTime) {
+    // Received a ping at the same priority level at which we're currently
+    // rendering. Restart from the root.
+    nextRoot = null;
   } else {
-    // Suspense already timed out. Compute a new expiration time
-    const currentTime = requestCurrentTime();
-    retryTime = computeExpirationForFiber(currentTime, boundaryFiber);
-    markPendingPriorityLevel(root, retryTime);
+    // Confirm that the root is still suspended at this level. Otherwise exit.
+    if (isPriorityLevelSuspended(root, pingTime)) {
+      // Ping at the original level
+      markPingedPriorityLevel(root, pingTime);
+      const rootExpirationTime = root.expirationTime;
+      if (rootExpirationTime !== NoWork) {
+        requestWork(root, rootExpirationTime);
+      }
+    }
   }
+}
 
-  // TODO: If the suspense fiber has already rendered the primary children
-  // without suspending (that is, all of the promises have already resolved),
-  // we should not trigger another update here. One case this happens is when
-  // we are in sync mode and a single promise is thrown both on initial render
-  // and on update; we attach two .then(retrySuspendedRoot) callbacks and each
-  // one performs Sync work, rerendering the Suspense.
+function retryTimedOutBoundary(boundaryFiber: Fiber, thenable: Thenable) {
+  // The boundary fiber (a Suspense component) previously timed out and was
+  // rendered in its fallback state. One of the promises that suspended it has
+  // resolved, which means at least part of the tree was likely unblocked. Try
+  // rendering again, at a new expiration time.
 
-  if ((boundaryFiber.mode & ConcurrentMode) !== NoContext) {
-    if (root === nextRoot && nextRenderExpirationTime === suspendedTime) {
-      // Received a ping at the same priority level at which we're currently
-      // rendering. Restart from the root.
-      nextRoot = null;
-    }
+  const retryCache: WeakSet<Thenable> | Set<Thenable> | null =
+    boundaryFiber.stateNode;
+  if (retryCache !== null) {
+    // The thenable resolved, so we no longer need to memoize, because it will
+    // never be thrown again.
+    retryCache.delete(thenable);
   }
 
-  scheduleWorkToRoot(boundaryFiber, retryTime);
-  if ((boundaryFiber.mode & ConcurrentMode) === NoContext) {
-    // Outside of concurrent mode, we must schedule an update on the source
-    // fiber, too, since it already committed in an inconsistent state and
-    // therefore does not have any pending work.
-    scheduleWorkToRoot(sourceFiber, retryTime);
-    const sourceTag = sourceFiber.tag;
-    if (sourceTag === ClassComponent && sourceFiber.stateNode !== null) {
-      // When we try rendering again, we should not reuse the current fiber,
-      // since it's known to be in an inconsistent state. Use a force updte to
-      // prevent a bail out.
-      const update = createUpdate(retryTime);
-      update.tag = ForceUpdate;
-      enqueueUpdate(sourceFiber, update);
+  const currentTime = requestCurrentTime();
+  const retryTime = computeExpirationForFiber(currentTime, boundaryFiber);
+  const root = scheduleWorkToRoot(boundaryFiber, retryTime);
+  if (root !== null) {
+    markPendingPriorityLevel(root, retryTime);
+    const rootExpirationTime = root.expirationTime;
+    if (rootExpirationTime !== NoWork) {
+      requestWork(root, rootExpirationTime);
     }
   }
-
-  const rootExpirationTime = root.expirationTime;
-  if (rootExpirationTime !== NoWork) {
-    requestWork(root, rootExpirationTime);
-  }
 }
 
 function scheduleWorkToRoot(fiber: Fiber, expirationTime): FiberRoot | null {
@@ -2550,7 +2545,8 @@ export {
   onUncaughtError,
   renderDidSuspend,
   renderDidError,
-  retrySuspendedRoot,
+  pingSuspendedRoot,
+  retryTimedOutBoundary,
   markLegacyErrorBoundaryAsFailed,
   isAlreadyFailedLegacyErrorBoundary,
   scheduleWork,

packages/react-reconciler/src/ReactFiberSuspenseComponent.js

@@ -14,10 +14,7 @@ export type SuspenseState = {|
   timedOutAt: ExpirationTime,
 |};
 
-export function shouldCaptureSuspense(
-  current: Fiber | null,
-  workInProgress: Fiber,
-): boolean {
+export function shouldCaptureSuspense(workInProgress: Fiber): boolean {
   // In order to capture, the Suspense component must have a fallback prop.
   if (workInProgress.memoizedProps.fallback === undefined) {
     return false;