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.

Author: acdlite

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__) {

packages/react/src/ReactStartTransition.js

@@ -23,12 +23,7 @@ export function startTransition(
   options?: StartTransitionOptions,
 ) {
   const prevTransition = ReactSharedInternals.T;
-  // Each renderer registers a callback to receive the return value of
-  // the scope function. This is used to implement async actions.
-  const callbacks = new Set<(BatchConfigTransition, mixed) => mixed>();
-  const transition: BatchConfigTransition = {
-    _callbacks: callbacks,
-  };
+  const transition: BatchConfigTransition = {};
   ReactSharedInternals.T = transition;
   const currentTransition = ReactSharedInternals.T;
 
@@ -48,12 +43,15 @@ export function startTransition(
   if (enableAsyncActions) {
     try {
       const returnValue = scope();
+      const onStartTransitionFinish = ReactSharedInternals.S;
+      if (onStartTransitionFinish !== null) {
+        onStartTransitionFinish(transition, returnValue);
+      }
       if (
         typeof returnValue === 'object' &&
         returnValue !== null &&
         typeof returnValue.then === 'function'
       ) {
-        callbacks.forEach(callback => callback(currentTransition, returnValue));
         returnValue.then(noop, reportGlobalError);
       }
     } catch (error) {