Fix memory leak after repeated setState bailouts (#25309)

acdlite · web-flow · commit d1bb1c586117 · 2022-09-22T11:19:54.000-04:00
There's a global queue (`concurrentQueues` in the ReactFiberConcurrentUpdates module) that is cleared at the beginning of each render phase. However, in the case of an eager `setState` bailout where the state is updated to same value as the current one, we add the update to the queue without scheduling a render. So the render phase never removes it from the queue. This can lead to a memory leak if it happens repeatedly without any other updates. There's only one place where this ever happens, so the fix was pretty straightforward. Currently there's no great way to test this from a Jest test, so I confirmed locally by checking in an existing test whether the array gets reset. @sompylasar had an interesting suggestion for how to catch these in the future: in the development build (perhaps behind a flag), use a Babel plugin to instrument all module-level variables. Then periodically sweep to confirm if something has leaked. The logic is that if there's no React work scheduled, and a module-level variable points to an object, it very likely indicates a memory leak.
diff --git a/packages/react-reconciler/src/ReactFiberConcurrentUpdates.new.js b/packages/react-reconciler/src/ReactFiberConcurrentUpdates.new.js
@@ -32,6 +32,7 @@ import {
 import {NoFlags, Placement, Hydrating} from './ReactFiberFlags';
 import {HostRoot, OffscreenComponent} from './ReactWorkTags';
 import {OffscreenVisible} from './ReactFiberOffscreenComponent';
+import {getWorkInProgressRoot} from './ReactFiberWorkLoop.new';
 
 export type ConcurrentUpdate = {
   next: ConcurrentUpdate,
@@ -139,6 +140,18 @@ export function enqueueConcurrentHookUpdateAndEagerlyBailout<S, A>(
   const concurrentQueue: ConcurrentQueue = (queue: any);
   const concurrentUpdate: ConcurrentUpdate = (update: any);
   enqueueUpdate(fiber, concurrentQueue, concurrentUpdate, lane);
+
+  // Usually we can rely on the upcoming render phase to process the concurrent
+  // queue. However, since this is a bail out, we're not scheduling any work
+  // here. So the update we just queued will leak until something else happens
+  // to schedule work (if ever).
+  //
+  // Check if we're currently in the middle of rendering a tree, and if not,
+  // process the queue immediately to prevent a leak.
+  const isConcurrentlyRendering = getWorkInProgressRoot() !== null;
+  if (!isConcurrentlyRendering) {
+    finishQueueingConcurrentUpdates();
+  }
 }
 
 export function enqueueConcurrentClassUpdate<State>(
diff --git a/packages/react-reconciler/src/ReactFiberConcurrentUpdates.old.js b/packages/react-reconciler/src/ReactFiberConcurrentUpdates.old.js
@@ -32,6 +32,7 @@ import {
 import {NoFlags, Placement, Hydrating} from './ReactFiberFlags';
 import {HostRoot, OffscreenComponent} from './ReactWorkTags';
 import {OffscreenVisible} from './ReactFiberOffscreenComponent';
+import {getWorkInProgressRoot} from './ReactFiberWorkLoop.old';
 
 export type ConcurrentUpdate = {
   next: ConcurrentUpdate,
@@ -139,6 +140,18 @@ export function enqueueConcurrentHookUpdateAndEagerlyBailout<S, A>(
   const concurrentQueue: ConcurrentQueue = (queue: any);
   const concurrentUpdate: ConcurrentUpdate = (update: any);
   enqueueUpdate(fiber, concurrentQueue, concurrentUpdate, lane);
+
+  // Usually we can rely on the upcoming render phase to process the concurrent
+  // queue. However, since this is a bail out, we're not scheduling any work
+  // here. So the update we just queued will leak until something else happens
+  // to schedule work (if ever).
+  //
+  // Check if we're currently in the middle of rendering a tree, and if not,
+  // process the queue immediately to prevent a leak.
+  const isConcurrentlyRendering = getWorkInProgressRoot() !== null;
+  if (!isConcurrentlyRendering) {
+    finishQueueingConcurrentUpdates();
+  }
 }
 
 export function enqueueConcurrentClassUpdate<State>(
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
@@ -1912,6 +1912,9 @@ function renderRootSync(root: FiberRoot, lanes: Lanes) {
   workInProgressRoot = null;
   workInProgressRootRenderLanes = NoLanes;
 
+  // It's safe to process the queue now that the render phase is complete.
+  finishQueueingConcurrentUpdates();
+
   return workInProgressRootExitStatus;
 }
 
@@ -2017,6 +2020,9 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
     workInProgressRoot = null;
     workInProgressRootRenderLanes = NoLanes;
 
+    // It's safe to process the queue now that the render phase is complete.
+    finishQueueingConcurrentUpdates();
+
     // Return the final exit status.
     return workInProgressRootExitStatus;
   }
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
@@ -1912,6 +1912,9 @@ function renderRootSync(root: FiberRoot, lanes: Lanes) {
   workInProgressRoot = null;
   workInProgressRootRenderLanes = NoLanes;
 
+  // It's safe to process the queue now that the render phase is complete.
+  finishQueueingConcurrentUpdates();
+
   return workInProgressRootExitStatus;
 }
 
@@ -2017,6 +2020,9 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
     workInProgressRoot = null;
     workInProgressRootRenderLanes = NoLanes;
 
+    // It's safe to process the queue now that the render phase is complete.
+    finishQueueingConcurrentUpdates();
+
     // Return the final exit status.
     return workInProgressRootExitStatus;
   }
