act: Move didScheduleLegacyUpdate to ensureRootIsScheduled (#26552)

`act` uses the `didScheduleLegacyUpdate` field to simulate the behavior
of batching in React <17 and below. It's a quirk leftover from a
previous implementation, not intentionally designed.

This sets `didScheduleLegacyUpdate` every time a legacy root receives an
update as opposed to only when the `executionContext` is empty. There's
no real reason to do it this way over some other way except that it's
how it used to work before #26512 and we should try our best to maintain
the existing behavior, quirks and all, since existing tests may have
come to accidentally rely on it.

This should fix some (though not all) of the internal Meta tests that
started failing after #26512 landed.

Will add a regression test before merging.

packages/react-reconciler/src/ReactFiberRootScheduler.js

@@ -119,6 +119,15 @@ export function ensureRootIsScheduled(root: FiberRoot): void {
     // unblock additional features we have planned.
     scheduleTaskForRootDuringMicrotask(root, now());
   }
+
+  if (
+    __DEV__ &&
+    ReactCurrentActQueue.isBatchingLegacy &&
+    root.tag === LegacyRoot
+  ) {
+    // Special `act` case: Record whenever a legacy update is scheduled.
+    ReactCurrentActQueue.didScheduleLegacyUpdate = true;
+  }
 }
 
 export function flushSyncWorkOnAllRoots() {

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -828,7 +828,6 @@ export function scheduleUpdateOnFiber(
     ) {
       if (__DEV__ && ReactCurrentActQueue.isBatchingLegacy) {
         // Treat `act` as if it's inside `batchedUpdates`, even in legacy mode.
-        ReactCurrentActQueue.didScheduleLegacyUpdate = true;
       } else {
         // Flush the synchronous work now, unless we're already working or inside
         // a batch. This is intentionally inside scheduleUpdateOnFiber instead of

packages/react-reconciler/src/__tests__/ReactIsomorphicAct-test.js

@@ -17,10 +17,13 @@ let Suspense;
 let DiscreteEventPriority;
 let startTransition;
 let waitForMicrotasks;
+let Scheduler;
+let assertLog;
 
 describe('isomorphic act()', () => {
   beforeEach(() => {
     React = require('react');
+    Scheduler = require('scheduler');
 
     ReactNoop = require('react-noop-renderer');
     DiscreteEventPriority =
@@ -31,6 +34,7 @@ describe('isomorphic act()', () => {
     startTransition = React.startTransition;
 
     waitForMicrotasks = require('internal-test-utils').waitForMicrotasks;
+    assertLog = require('internal-test-utils').assertLog;
   });
 
   beforeEach(() => {
@@ -41,6 +45,11 @@ describe('isomorphic act()', () => {
     jest.restoreAllMocks();
   });
 
+  function Text({text}) {
+    Scheduler.log(text);
+    return text;
+  }
+
   // @gate __DEV__
   test('bypasses queueMicrotask', async () => {
     const root = ReactNoop.createRoot();
@@ -132,19 +141,67 @@ describe('isomorphic act()', () => {
     const root = ReactNoop.createLegacyRoot();
 
     await act(async () => {
-      // These updates are batched. This replicates the behavior of the original
-      // `act` implementation, for compatibility.
-      root.render('A');
-      root.render('B');
-      // Nothing has rendered yet.
-      expect(root).toMatchRenderedOutput(null);
-      await null;
-      // Updates are flushed after the first await.
-      expect(root).toMatchRenderedOutput('B');
+      queueMicrotask(() => {
+        Scheduler.log('Current tree in microtask: ' + root.getChildrenAsJSX());
+        root.render(<Text text="C" />);
+      });
+      root.render(<Text text="A" />);
+      root.render(<Text text="B" />);
 
-      // Subsequent updates in the same scope aren't batched.
-      root.render('C');
-      expect(root).toMatchRenderedOutput('C');
+      await null;
+      assertLog([
+        // A and B should render in a single batch _before_ the microtask queue
+        // has run. This replicates the behavior of the original `act`
+        // implementation, for compatibility.
+        'B',
+        'Current tree in microtask: B',
+
+        // C isn't scheduled until a microtask, so it's rendered separately.
+        'C',
+      ]);
+
+      // Subsequent updates should also render in separate batches.
+      root.render(<Text text="D" />);
+      root.render(<Text text="E" />);
+      assertLog(['D', 'E']);
+    });
+  });
+
+  // @gate __DEV__
+  test('in legacy mode, in an async scope, updates are batched until the first `await` (regression test: batchedUpdates)', async () => {
+    const root = ReactNoop.createLegacyRoot();
+
+    await act(async () => {
+      queueMicrotask(() => {
+        Scheduler.log('Current tree in microtask: ' + root.getChildrenAsJSX());
+        root.render(<Text text="C" />);
+      });
+
+      // This is a regression test. The presence of `batchedUpdates` would cause
+      // these updates to not flush until a microtask. The correct behavior is
+      // that they flush before the microtask queue, regardless of whether
+      // they are wrapped with `batchedUpdates`.
+      ReactNoop.batchedUpdates(() => {
+        root.render(<Text text="A" />);
+        root.render(<Text text="B" />);
+      });
+
+      await null;
+      assertLog([
+        // A and B should render in a single batch _before_ the microtask queue
+        // has run. This replicates the behavior of the original `act`
+        // implementation, for compatibility.
+        'B',
+        'Current tree in microtask: B',
+
+        // C isn't scheduled until a microtask, so it's rendered separately.
+        'C',
+      ]);
+
+      // Subsequent updates should also render in separate batches.
+      root.render(<Text text="D" />);
+      root.render(<Text text="E" />);
+      assertLog(['D', 'E']);
     });
   });
 

scripts/jest/matchers/reactTestMatchers.js

@@ -20,13 +20,13 @@ function captureAssertion(fn) {
   return {pass: true};
 }
 
-function assertYieldsWereCleared(Scheduler) {
+function assertYieldsWereCleared(Scheduler, caller) {
   const actualYields = Scheduler.unstable_clearLog();
   if (actualYields.length !== 0) {
     const error = Error(
       'The event log is not empty. Call assertLog(...) first.'
     );
-    Error.captureStackTrace(error, assertYieldsWereCleared);
+    Error.captureStackTrace(error, caller);
     throw error;
   }
 }
@@ -34,7 +34,7 @@ function assertYieldsWereCleared(Scheduler) {
 function toMatchRenderedOutput(ReactNoop, expectedJSX) {
   if (typeof ReactNoop.getChildrenAsJSX === 'function') {
     const Scheduler = ReactNoop._Scheduler;
-    assertYieldsWereCleared(Scheduler);
+    assertYieldsWereCleared(Scheduler, toMatchRenderedOutput);
     return captureAssertion(() => {
       expect(ReactNoop.getChildrenAsJSX()).toEqual(expectedJSX);
     });