act: Move didScheduleLegacyUpdate to ensureRootIsScheduled

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);
+      queueMicrotask(() => {
+        Scheduler.log('Current tree in microtask: ' + root.getChildrenAsJSX());
+        root.render(<Text text="C" />);
+      });
+      root.render(<Text text="A" />);
+      root.render(<Text text="B" />);
+
       await null;
-      // Updates are flushed after the first await.
-      expect(root).toMatchRenderedOutput('B');
+      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']);
+    });
+  });
 
-      // Subsequent updates in the same scope aren't batched.
-      root.render('C');
-      expect(root).toMatchRenderedOutput('C');
+  // @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,21 +20,21 @@ 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;
   }
 }
 
 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);
     });