From 49f7410467950de57bd1d9509e67de91185f7bad Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Sun, 5 Mar 2023 20:24:39 -0500
Subject: [PATCH] Fix: Infinite act loop caused by wrong shouldYield (#26317)

Based on a bug report from @bvaughn.

`act` should not consult `shouldYield` when it's performing work,
because in a unit testing environment, I/O (such as `setTimeout`) is
likely mocked. So the result of `shouldYield` can't be trusted.

In the regression test, I simulate the bug by mocking `shouldYield` to
always return `true`. This causes an infinite loop in `act`, because it
will keep trying to render and React will keep yielding.
---
 .../src/ReactFiberWorkLoop.js                 | 12 +++-
 .../ReactSchedulerIntegration-test.js         | 68 +++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.js b/packages/react-reconciler/src/ReactFiberWorkLoop.js
index fd6e51a0872bc..1b34601740f60 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.js
+++ b/packages/react-reconciler/src/ReactFiberWorkLoop.js
@@ -2268,7 +2268,17 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
           }
         }
       }
-      workLoopConcurrent();
+
+      if (__DEV__ && ReactCurrentActQueue.current !== null) {
+        // `act` special case: If we're inside an `act` scope, don't consult
+        // `shouldYield`. Always keep working until the render is complete.
+        // This is not just an optimization: in a unit test environment, we
+        // can't trust the result of `shouldYield`, because the host I/O is
+        // likely mocked.
+        workLoopSync();
+      } else {
+        workLoopConcurrent();
+      }
       break;
     } catch (thrownValue) {
       handleThrow(root, thrownValue);
diff --git a/packages/react-reconciler/src/__tests__/ReactSchedulerIntegration-test.js b/packages/react-reconciler/src/__tests__/ReactSchedulerIntegration-test.js
index 1dffc1be85839..f3b37d4e3460c 100644
--- a/packages/react-reconciler/src/__tests__/ReactSchedulerIntegration-test.js
+++ b/packages/react-reconciler/src/__tests__/ReactSchedulerIntegration-test.js
@@ -320,3 +320,71 @@ describe(
     });
   },
 );
+
+describe('`act` bypasses Scheduler methods completely,', () => {
+  let infiniteLoopGuard;
+
+  beforeEach(() => {
+    jest.resetModules();
+
+    infiniteLoopGuard = 0;
+
+    jest.mock('scheduler', () => {
+      const actual = jest.requireActual('scheduler/unstable_mock');
+      return {
+        ...actual,
+        unstable_shouldYield() {
+          // This simulates a bug report where `shouldYield` returns true in a
+          // unit testing environment. Because `act` will keep working until
+          // there's no more work left, it would fall into an infinite loop.
+          // The fix is that when performing work inside `act`, we should bypass
+          // `shouldYield` completely, because we can't trust it to be correct.
+          if (infiniteLoopGuard++ > 100) {
+            throw new Error('Detected an infinite loop');
+          }
+          return true;
+        },
+      };
+    });
+
+    React = require('react');
+    ReactNoop = require('react-noop-renderer');
+    startTransition = React.startTransition;
+  });
+
+  afterEach(() => {
+    jest.mock('scheduler', () => jest.requireActual('scheduler/unstable_mock'));
+  });
+
+  // @gate __DEV__
+  it('inside `act`, does not call `shouldYield`, even during a concurrent render', async () => {
+    function App() {
+      return (
+        <>
+          <div>A</div>
+          <div>B</div>
+          <div>C</div>
+        </>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+    const publicAct = React.unstable_act;
+    const prevIsReactActEnvironment = global.IS_REACT_ACT_ENVIRONMENT;
+    try {
+      global.IS_REACT_ACT_ENVIRONMENT = true;
+      await publicAct(async () => {
+        startTransition(() => root.render(<App />));
+      });
+    } finally {
+      global.IS_REACT_ACT_ENVIRONMENT = prevIsReactActEnvironment;
+    }
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>A</div>
+        <div>B</div>
+        <div>C</div>
+      </>,
+    );
+  });
+});