Fix async batching in React.startTransition

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,48 @@ 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(fn) {
+//     return startTransitionDOM(() => {
+//       return startTransitionART(() => {
+//         return startTransitionThreeFiber(() => {
+//           // and so on...
+//           return fn();
+//         });
+//       });
+//     });
+//   }
+//
+// 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__) {