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

Author: acdlite

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';
 importtype{CapturedValue,CapturedError}from'./ReactCapturedValue';
 importtype{SuspenseState}from'./ReactFiberSuspenseComponent';
 importtype{FunctionComponentUpdateQueue}from'./ReactFiberHooks';
+importtype{Thenable}from'./ReactFiberScheduler';
 
+import{unstable_wrapasSchedule_tracing_wrap}from'scheduler/tracing';
 import{
 enableHooks,
 enableSchedulerTracing,
@@ -88,6 +90,7 @@ import {
 import{
 captureCommitPhaseError,
 requestCurrentTime,
+retryTimedOutBoundary,
 }from'./ReactFiberScheduler';
 import{
 NoEffectasNoHookEffect,
@@ -106,6 +109,8 @@ if (__DEV__) {
 didWarnAboutUndefinedSnapshotBeforeUpdate=newSet();
 }
 
+constPossiblyWeakSet=typeofWeakSet==='function' ? WeakSet : Set;
+
 exportfunctionlogError(boundary: Fiber,errorInfo: CapturedValue<mixed>){
 constsource=errorInfo.source;
 letstack=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.
+constthenables: Set<Thenable> | null = (finishedWork.updateQueue: any);
+ if (thenables !== null) {
+finishedWork.updateQueue=null;
+letretryCache=finishedWork.stateNode;
+if(retryCache===null){
+retryCache=finishedWork.stateNode=newPossiblyWeakSet();
+}
+thenables.forEach(thenable=>{
+// Memoize using the boundary fiber to prevent redundant listeners.
+letretry=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.
 constlatestPendingTime=root.latestPendingTime;
 if(latestPendingTime!==NoWork){
@@ -209,10 +213,8 @@ export function markPingedPriorityLevel(
 }
 
 functionclearPing(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.
 constlatestPingedTime=root.latestPingedTime;
-if(latestPingedTime!==NoWork&&latestPingedTime>=completedTime){
+if(latestPingedTime>=completedTime){
 root.latestPingedTime=NoWork;
 }
 }

packages/react-reconciler/src/ReactFiberRoot.js

@@ -10,6 +10,7 @@
 importtype{Fiber}from'./ReactFiber';
 importtype{ExpirationTime}from'./ReactFiberExpirationTime';
 importtype{TimeoutHandle,NoTimeout}from'./ReactFiberHostConfig';
+importtype{Thenable}from'./ReactFiberScheduler';
 importtype{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{
 isContextProviderasisLegacyContextProvider,
@@ -1646,62 +1641,62 @@ function renderDidError() {
 nextRenderDidError=true;
 }
 
-functionretrySuspendedRoot(
+functionpingSuspendedRoot(
 root: FiberRoot,
-boundaryFiber: Fiber,
-sourceFiber: Fiber,
-suspendedTime: ExpirationTime,
+thenable: Thenable,
+pingTime: ExpirationTime,
 ){
-letretryTime;
+// 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;
+constpingCache=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
-constcurrentTime=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);
+constrootExpirationTime=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.
+functionretryTimedOutBoundary(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;
-}
+constretryCache: 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);
-constsourceTag=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.
-constupdate=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);
+constrootExpirationTime=root.expirationTime;
+if(rootExpirationTime!==NoWork){
+requestWork(root,rootExpirationTime);
 }
 }
-
-constrootExpirationTime=root.expirationTime;
-if(rootExpirationTime!==NoWork){
-requestWork(root,rootExpirationTime);
-}
 }
 
 functionscheduleWorkToRoot(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,
 |};
 
-exportfunctionshouldCaptureSuspense(
-current: Fiber|null,
-workInProgress: Fiber,
-): boolean{
+exportfunctionshouldCaptureSuspense(workInProgress: Fiber): boolean{
 // In order to capture, the Suspense component must have a fallback prop.
 if(workInProgress.memoizedProps.fallback===undefined){
 returnfalse;