act: Batch updates, even in legacy roots

In legacy roots, if an update originates outside of `batchedUpdates`,
check if it's inside an `act` scope; if so, treat it as if it were
batched. This is only necessary in legacy roots because in concurrent
roots, updates are batched by default.

With this change, the Test Utils and Test Renderer versions of `act` are
nothing more than aliases of the isomorphic API (still not exposed, but
will likely be the recommended API that replaces the others).

packages/react-devtools-shared/src/__tests__/inspectedElement-test.js

@@ -2462,7 +2462,7 @@ describe('InspectedElement', () => {
       };
       const toggleError = async forceError => {
         await withErrorsOrWarningsIgnored(['ErrorBoundary'], async () => {
-          await utils.actAsync(() => {
+          await TestUtilsAct(() => {
             bridge.send('overrideError', {
               id: targetErrorBoundaryID,
               rendererID: store.getRendererIDForElement(targetErrorBoundaryID),

packages/react-dom/src/test-utils/ReactTestUtils.js

@@ -33,14 +33,8 @@ const getNodeFromInstance = EventInternals[1];
 const getFiberCurrentPropsFromNode = EventInternals[2];
 const enqueueStateRestore = EventInternals[3];
 const restoreStateIfNeeded = EventInternals[4];
-const batchedUpdates = EventInternals[5];
 
-const act_notBatchedInLegacyMode = React.unstable_act;
-function act(callback) {
-  return act_notBatchedInLegacyMode(() => {
-    return batchedUpdates(callback);
-  });
-}
+const act = React.unstable_act;
 
 function Event(suffix) {}
 

packages/react-reconciler/src/ReactFiberWorkLoop.new.js

@@ -518,7 +518,9 @@ export function scheduleUpdateOnFiber(
   if (
     lane === SyncLane &&
     executionContext === NoContext &&
-    (fiber.mode & ConcurrentMode) === NoMode
+    (fiber.mode & ConcurrentMode) === NoMode &&
+    // Treat `act` as if it's inside `batchedUpdates`, even in legacy mode.
+    !(__DEV__ && ReactCurrentActQueue.isBatchingLegacy)
   ) {
     // Flush the synchronous work now, unless we're already working or inside
     // a batch. This is intentionally inside scheduleUpdateOnFiber instead of
@@ -668,6 +670,9 @@ function ensureRootIsScheduled(root: FiberRoot, currentTime: number) {
     // Special case: Sync React callbacks are scheduled on a special
     // internal queue
     if (root.tag === LegacyRoot) {
+      if (__DEV__ && ReactCurrentActQueue.isBatchingLegacy !== null) {
+        ReactCurrentActQueue.didScheduleLegacyUpdate = true;
+      }
       scheduleLegacySyncCallback(performSyncWorkOnRoot.bind(null, root));
     } else {
       scheduleSyncCallback(performSyncWorkOnRoot.bind(null, root));
@@ -1049,7 +1054,11 @@ export function batchedUpdates<A, R>(fn: A => R, a: A): R {
     executionContext = prevExecutionContext;
     // If there were legacy sync updates, flush them at the end of the outer
     // most batchedUpdates-like method.
-    if (executionContext === NoContext) {
+    if (
+      executionContext === NoContext &&
+      // Treat `act` as if it's inside `batchedUpdates`, even in legacy mode.
+      !(__DEV__ && ReactCurrentActQueue.isBatchingLegacy)
+    ) {
       resetRenderTimer();
       flushSyncCallbacksOnlyInLegacyMode();
     }

packages/react-reconciler/src/ReactFiberWorkLoop.old.js

@@ -518,7 +518,9 @@ export function scheduleUpdateOnFiber(
   if (
     lane === SyncLane &&
     executionContext === NoContext &&
-    (fiber.mode & ConcurrentMode) === NoMode
+    (fiber.mode & ConcurrentMode) === NoMode &&
+    // Treat `act` as if it's inside `batchedUpdates`, even in legacy mode.
+    !(__DEV__ && ReactCurrentActQueue.isBatchingLegacy)
   ) {
     // Flush the synchronous work now, unless we're already working or inside
     // a batch. This is intentionally inside scheduleUpdateOnFiber instead of
@@ -668,6 +670,9 @@ function ensureRootIsScheduled(root: FiberRoot, currentTime: number) {
     // Special case: Sync React callbacks are scheduled on a special
     // internal queue
     if (root.tag === LegacyRoot) {
+      if (__DEV__ && ReactCurrentActQueue.isBatchingLegacy !== null) {
+        ReactCurrentActQueue.didScheduleLegacyUpdate = true;
+      }
       scheduleLegacySyncCallback(performSyncWorkOnRoot.bind(null, root));
     } else {
       scheduleSyncCallback(performSyncWorkOnRoot.bind(null, root));
@@ -1049,7 +1054,11 @@ export function batchedUpdates<A, R>(fn: A => R, a: A): R {
     executionContext = prevExecutionContext;
     // If there were legacy sync updates, flush them at the end of the outer
     // most batchedUpdates-like method.
-    if (executionContext === NoContext) {
+    if (
+      executionContext === NoContext &&
+      // Treat `act` as if it's inside `batchedUpdates`, even in legacy mode.
+      !(__DEV__ && ReactCurrentActQueue.isBatchingLegacy)
+    ) {
       resetRenderTimer();
       flushSyncCallbacksOnlyInLegacyMode();
     }

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

@@ -78,4 +78,53 @@ describe('isomorphic act()', () => {
     });
     expect(returnValue).toEqual('hi');
   });
+
+  // @gate __DEV__
+  test('in legacy mode, updates are batched', () => {
+    const root = ReactNoop.createLegacyRoot();
+
+    // Outside of `act`, legacy updates are flushed completely synchronously
+    root.render('A');
+    expect(root).toMatchRenderedOutput('A');
+
+    // `act` will batch the updates and flush them at the end
+    act(() => {
+      root.render('B');
+      // Hasn't flushed yet
+      expect(root).toMatchRenderedOutput('A');
+
+      // Confirm that a nested `batchedUpdates` call won't cause the updates
+      // to flush early.
+      ReactNoop.batchedUpdates(() => {
+        root.render('C');
+      });
+
+      // Still hasn't flushed
+      expect(root).toMatchRenderedOutput('A');
+    });
+
+    // Now everything renders in a single batch.
+    expect(root).toMatchRenderedOutput('C');
+  });
+
+  // @gate __DEV__
+  test('in legacy mode, in an async scope, updates are batched until the first `await`', async () => {
+    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');
+
+      // Subsequent updates in the same scope aren't batched.
+      root.render('C');
+      expect(root).toMatchRenderedOutput('C');
+    });
+  });
 });

packages/react-test-renderer/src/ReactTestRenderer.js

@@ -7,7 +7,6 @@
  * @flow
  */
 
-import type {Thenable} from 'shared/ReactTypes';
 import type {Fiber} from 'react-reconciler/src/ReactInternalTypes';
 import type {FiberRoot} from 'react-reconciler/src/ReactInternalTypes';
 import type {Instance, TextInstance} from './ReactTestHostConfig';
@@ -50,12 +49,7 @@ import {getPublicInstance} from './ReactTestHostConfig';
 import {ConcurrentRoot, LegacyRoot} from 'react-reconciler/src/ReactRootTags';
 import {allowConcurrentByDefault} from 'shared/ReactFeatureFlags';
 
-const act_notBatchedInLegacyMode = React.unstable_act;
-function act<T>(callback: () => T): Thenable<T> {
-  return act_notBatchedInLegacyMode(() => {
-    return batchedUpdates(callback);
-  });
-}
+const act = React.unstable_act;
 
 // TODO: Remove from public bundle
 

packages/react/src/ReactAct.js

@@ -28,12 +28,34 @@ export function act<T>(callback: () => T | Thenable<T>): Thenable<T> {
       ReactCurrentActQueue.current = [];
     }
 
+    const prevIsBatchingLegacy = ReactCurrentActQueue.isBatchingLegacy;
     let result;
     try {
+      // Used to reproduce behavior of `batchedUpdates` in legacy mode. Only
+      // set to `true` while the given callback is executed, not for updates
+      // triggered during an async event, because this is how the legacy
+      // implementation of `act` behaved.
+      ReactCurrentActQueue.isBatchingLegacy = true;
       result = callback();
+
+      // Replicate behavior of original `act` implementation in legacy mode,
+      // which flushed updates immediately after the scope function exits, even
+      // if it's an async function.
+      if (
+        !prevIsBatchingLegacy &&
+        ReactCurrentActQueue.didScheduleLegacyUpdate
+      ) {
+        const queue = ReactCurrentActQueue.current;
+        if (queue !== null) {
+          ReactCurrentActQueue.didScheduleLegacyUpdate = false;
+          flushActQueue(queue);
+        }
+      }
     } catch (error) {
       popActScope(prevActScopeDepth);
       throw error;
+    } finally {
+      ReactCurrentActQueue.isBatchingLegacy = prevIsBatchingLegacy;
     }
 
     if (

packages/react/src/ReactCurrentActQueue.js

@@ -17,6 +17,10 @@ const ReactCurrentActQueue = {
   // on at the testing frameworks layer? Instead of what we do now, which
   // is check if a `jest` global is defined.
   disableActWarning: (false: boolean),
+
+  // Used to reproduce behavior of `batchedUpdates` in legacy mode.
+  isBatchingLegacy: false,
+  didScheduleLegacyUpdate: false,
 };
 
 export default ReactCurrentActQueue;