Fix async batching in React.startTransition

Discovered a flaw in the current implementation of the isomorphic
startTransition implementation (React.startTransition), related to async
actions. It only creates an async scope if something calls setState
within the synchronous part of the action (i.e. before the first
`await`). I had thought this was fine because if there's no update
during this part, then there's nothing that needs to be entangled. I
didn't think this through, though — if there are multiple async updates
interleaved throughout the rest of the action, we need the async scope
to have already been created so that _those_ are batched together.

An even more easier way to observe this is to schedule an optimistic
update after an `await` — the optimistic update should not be reverted
until the async action is complete.

To implement, during the reconciler's module initialization, we compose
its startTransition implementation with any previous reconciler's
startTransition that was already initialized. Then, the isomorphic
startTransition is the composition of every
reconciler's startTransition.

function startTransition() {
  startTransitionDOM(() => {
    startTransitionART(() => {
      startTransitionThreeFiber(() => {
        // and so on...
      });
    });
  });
});

This is basically how flushSync is implemented, too.

packages/react-dom/src/__tests__/ReactStartTransitionMultipleRenderers-test.js

@@ -0,0 +1,143 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ */
+
+'use strict';
+
+describe('ReactStartTransitionMultipleRenderers', () => {
+  let act;
+  let container;
+  let React;
+  let ReactDOMClient;
+  let Scheduler;
+  let assertLog;
+  let startTransition;
+  let useOptimistic;
+  let textCache;
+
+  beforeEach(() => {
+    jest.resetModules();
+    React = require('react');
+    ReactDOMClient = require('react-dom/client');
+    Scheduler = require('scheduler');
+    act = require('internal-test-utils').act;
+    assertLog = require('internal-test-utils').assertLog;
+    startTransition = React.startTransition;
+    useOptimistic = React.useOptimistic;
+    container = document.createElement('div');
+    document.body.appendChild(container);
+
+    textCache = new Map();
+  });
+
+  function resolveText(text) {
+    const record = textCache.get(text);
+    if (record === undefined) {
+      const newRecord = {
+        status: 'resolved',
+        value: text,
+      };
+      textCache.set(text, newRecord);
+    } else if (record.status === 'pending') {
+      const thenable = record.value;
+      record.status = 'resolved';
+      record.value = text;
+      thenable.pings.forEach(t => t(text));
+    }
+  }
+
+  function getText(text) {
+    const record = textCache.get(text);
+    if (record === undefined) {
+      const thenable = {
+        pings: [],
+        then(resolve) {
+          if (newRecord.status === 'pending') {
+            thenable.pings.push(resolve);
+          } else {
+            Promise.resolve().then(() => resolve(newRecord.value));
+          }
+        },
+      };
+      const newRecord = {
+        status: 'pending',
+        value: thenable,
+      };
+      textCache.set(text, newRecord);
+      return thenable;
+    } else {
+      switch (record.status) {
+        case 'pending':
+          return record.value;
+        case 'rejected':
+          return Promise.reject(record.value);
+        case 'resolved':
+          return Promise.resolve(record.value);
+      }
+    }
+  }
+
+  function Text({text}) {
+    Scheduler.log(text);
+    return text;
+  }
+
+  afterEach(() => {
+    document.body.removeChild(container);
+  });
+
+  // This test imports multiple reconcilers. Because of how the renderers are
+  // built, it only works when running tests using the actual build artifacts,
+  // not the source files.
+  // @gate !source
+  test('React.startTransition works across multiple renderers', async () => {
+    const ReactNoop = require('react-noop-renderer');
+
+    const setIsPendings = new Set();
+
+    function App() {
+      const [isPending, setIsPending] = useOptimistic(false);
+      setIsPendings.add(setIsPending);
+      return <Text text={isPending ? 'Pending' : 'Not pending'} />;
+    }
+
+    const noopRoot = ReactNoop.createRoot(null);
+    const domRoot = ReactDOMClient.createRoot(container);
+
+    // Run the same component in two separate renderers.
+    await act(() => {
+      noopRoot.render(<App />);
+      domRoot.render(<App />);
+    });
+    assertLog(['Not pending', 'Not pending']);
+    expect(container.textContent).toEqual('Not pending');
+    expect(noopRoot).toMatchRenderedOutput('Not pending');
+
+    await act(() => {
+      startTransition(async () => {
+        // Wait until after an async gap before setting the optimistic state.
+        await getText('Wait before setting isPending');
+        setIsPendings.forEach(setIsPending => setIsPending(true));
+
+        // The optimistic state should not be reverted until the
+        // action completes.
+        await getText('Wait until end of async action');
+      });
+    });
+
+    await act(() => resolveText('Wait before setting isPending'));
+    assertLog(['Pending', 'Pending']);
+    expect(container.textContent).toEqual('Pending');
+    expect(noopRoot).toMatchRenderedOutput('Pending');
+
+    await act(() => resolveText('Wait until end of async action'));
+    assertLog(['Not pending', 'Not pending']);
+    expect(container.textContent).toEqual('Not pending');
+    expect(noopRoot).toMatchRenderedOutput('Not pending');
+  });
+});

packages/react-reconciler/src/ReactFiberHooks.js

@@ -154,10 +154,7 @@ import {
 import {HostTransitionContext} from './ReactFiberHostContext';
 import {requestTransitionLane} from './ReactFiberRootScheduler';
 import {isCurrentTreeHidden} from './ReactFiberHiddenContext';
-import {
-  notifyTransitionCallbacks,
-  requestCurrentTransition,
-} from './ReactFiberTransition';
+import {requestCurrentTransition} from './ReactFiberTransition';
 
 export type Update<S, A> = {
   lane: Lane,
@@ -2020,9 +2017,7 @@ function runActionStateAction<S, P>(
 
   // This is a fork of startTransition
   const prevTransition = ReactSharedInternals.T;
-  const currentTransition: BatchConfigTransition = {
-    _callbacks: new Set<(BatchConfigTransition, mixed) => mixed>(),
-  };
+  const currentTransition: BatchConfigTransition = {};
   ReactSharedInternals.T = currentTransition;
   if (__DEV__) {
     ReactSharedInternals.T._updatedFibers = new Set();
@@ -2034,14 +2029,17 @@ function runActionStateAction<S, P>(
 
   try {
     const returnValue = action(prevState, payload);
+    const onStartTransitionFinish = ReactSharedInternals.S;
+    if (onStartTransitionFinish !== null) {
+      onStartTransitionFinish(currentTransition, returnValue);
+    }
     if (
       returnValue !== null &&
       typeof returnValue === 'object' &&
       // $FlowFixMe[method-unbinding]
       typeof returnValue.then === 'function'
     ) {
       const thenable = ((returnValue: any): Thenable<Awaited<S>>);
-      notifyTransitionCallbacks(currentTransition, thenable);
 
       // Attach a listener to read the return state of the action. As soon as
       // this resolves, we can run the next action in the sequence.
@@ -2843,9 +2841,7 @@ function startTransition<S>(
   );
 
   const prevTransition = ReactSharedInternals.T;
-  const currentTransition: BatchConfigTransition = {
-    _callbacks: new Set<(BatchConfigTransition, mixed) => mixed>(),
-  };
+  const currentTransition: BatchConfigTransition = {};
 
   if (enableAsyncActions) {
     // We don't really need to use an optimistic update here, because we
@@ -2876,6 +2872,10 @@ function startTransition<S>(
   try {
     if (enableAsyncActions) {
       const returnValue = callback();
+      const onStartTransitionFinish = ReactSharedInternals.S;
+      if (onStartTransitionFinish !== null) {
+        onStartTransitionFinish(currentTransition, returnValue);
+      }
 
       // Check if we're inside an async action scope. If so, we'll entangle
       // this new action with the existing scope.
@@ -2891,7 +2891,6 @@ function startTransition<S>(
         typeof returnValue.then === 'function'
       ) {
         const thenable = ((returnValue: any): Thenable<mixed>);
-        notifyTransitionCallbacks(currentTransition, thenable);
         // Create a thenable that resolves to `finishedState` once the async
         // action has completed.
         const thenableForFinishedState = chainThenableValue(

packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js

@@ -47,7 +47,6 @@ export type BatchConfigTransition = {
   name?: string,
   startTime?: number,
   _updatedFibers?: Set<Fiber>,
-  _callbacks: Set<(BatchConfigTransition, mixed) => mixed>,
 };
 
 // TODO: Is there a way to not include the tag or name here?

packages/react-reconciler/src/ReactFiberTransition.js

@@ -38,32 +38,47 @@ import {entangleAsyncAction} from './ReactFiberAsyncAction';
 
 export const NoTransition = null;
 
-export function requestCurrentTransition(): BatchConfigTransition | null {
-  const transition = ReactSharedInternals.T;
-  if (transition !== null) {
-    // Whenever a transition update is scheduled, register a callback on the
-    // transition object so we can get the return value of the scope function.
-    transition._callbacks.add((handleAsyncAction: any));
-  }
-  return transition;
-}
-
-function handleAsyncAction(
+// Attach this reconciler instance's onStartTransitionFinish implementation to
+// the shared internals object. This is used by the isomorphic implementation of
+// startTransition to compose all the startTransitions together.
+//
+//   function startTransition() {
+//     startTransitionDOM(() => {
+//       startTransitionART(() => {
+//         startTransitionThreeFiber(() => {
+//           // and so on...
+//         });
+//       });
+//     });
+//   })
+//
+// Currently we only compose together the code that runs at the end of each
+// startTransition, because for now that's sufficient — the part that sets
+// isTransition=true on the stack uses a separate shared internal field. But
+// really we should delete the shared field and track isTransition per
+// reconciler. Leaving this for a future PR.
+const prevOnStartTransitionFinish = ReactSharedInternals.S;
+ReactSharedInternals.S = function onStartTransitionFinishForReconciler(
   transition: BatchConfigTransition,
-  thenable: Thenable<mixed>,
-): void {
-  if (enableAsyncActions) {
-    // This is an async action.
+  returnValue: mixed,
+) {
+  if (
+    enableAsyncActions &&
+    typeof returnValue === 'object' &&
+    returnValue !== null &&
+    typeof returnValue.then === 'function'
+  ) {
+    // This is an async action
+    const thenable: Thenable<mixed> = (returnValue: any);
     entangleAsyncAction(transition, thenable);
   }
-}
+  if (prevOnStartTransitionFinish !== null) {
+    prevOnStartTransitionFinish(transition, returnValue);
+  }
+};
 
-export function notifyTransitionCallbacks(
-  transition: BatchConfigTransition,
-  returnValue: mixed,
-) {
-  const callbacks = transition._callbacks;
-  callbacks.forEach(callback => callback(transition, returnValue));
+export function requestCurrentTransition(): BatchConfigTransition | null {
+  return ReactSharedInternals.T;
 }
 
 // When retrying a Suspense/Offscreen boundary, we restore the cache that was

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

@@ -1731,6 +1731,76 @@ describe('ReactAsyncActions', () => {
     expect(root).toMatchRenderedOutput(<span>Updated</span>);
   });
 
+  // @gate enableAsyncActions
+  test(
+    'regression: updates in an action passed to React.startTransition are batched ' +
+      'even if there were no updates before the first await',
+    async () => {
+      // Regression for a bug that occured in an older, too-clever-by-half
+      // implementation of the isomorphic startTransition API. Now, the
+      // isomorphic startTransition is literally the composition of every
+      // reconciler instance's startTransition, so the behavior is less likely
+      // to regress in the future.
+      const startTransition = React.startTransition;
+
+      let setOptimisticText;
+      function App({text: canonicalText}) {
+        const [text, _setOptimisticText] = useOptimistic(
+          canonicalText,
+          (_, optimisticText) => `${optimisticText} (loading...)`,
+        );
+        setOptimisticText = _setOptimisticText;
+        return (
+          <span>
+            <Text text={text} />
+          </span>
+        );
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(() => {
+        root.render(<App text="Initial" />);
+      });
+      assertLog(['Initial']);
+      expect(root).toMatchRenderedOutput(<span>Initial</span>);
+
+      // Start an async action using the non-hook form of startTransition. The
+      // action includes an optimistic update.
+      await act(() => {
+        startTransition(async () => {
+          Scheduler.log('Async action started');
+
+          // Yield to an async task *before* any updates have occurred.
+          await getText('Yield before optimistic update');
+
+          // This optimistic update happens after an async gap. In the
+          // regression case, this update was not correctly associated with
+          // the outer async action, causing the optimistic update to be
+          // immediately reverted.
+          setOptimisticText('Updated');
+
+          await getText('Yield before updating');
+          Scheduler.log('Async action ended');
+          startTransition(() => root.render(<App text="Updated" />));
+        });
+      });
+      assertLog(['Async action started']);
+
+      // Wait for an async gap, then schedule an optimistic update.
+      await act(() => resolveText('Yield before optimistic update'));
+
+      // Because the action hasn't finished yet, the optimistic UI is shown.
+      assertLog(['Updated (loading...)']);
+      expect(root).toMatchRenderedOutput(<span>Updated (loading...)</span>);
+
+      // Finish the async action. The optimistic state is reverted and replaced
+      // by the canonical state.
+      await act(() => resolveText('Yield before updating'));
+      assertLog(['Async action ended', 'Updated']);
+      expect(root).toMatchRenderedOutput(<span>Updated</span>);
+    },
+  );
+
   test('React.startTransition captures async errors and passes them to reportError', async () => {
     // NOTE: This is gated here instead of using the pragma because the failure
     // happens asynchronously and the `gate` runtime doesn't capture it.

packages/react/src/ReactSharedInternalsClient.js

@@ -15,6 +15,7 @@ export type SharedStateClient = {
   H: null | Dispatcher, // ReactCurrentDispatcher for Hooks
   A: null | AsyncDispatcher, // ReactCurrentCache for Cache
   T: null | BatchConfigTransition, // ReactCurrentBatchConfig for Transitions
+  S: null | ((BatchConfigTransition, mixed) => void), // onStartTransitionFinish
 
   // DEV-only
 
@@ -45,6 +46,7 @@ const ReactSharedInternals: SharedStateClient = ({
   H: null,
   A: null,
   T: null,
+  S: null,
 }: any);
 
 if (__DEV__) {