Fix async batching in React.startTransition (#29226)

Note: Despite the similar-sounding description, this fix is unrelated to
the issue where updates that occur after an `await` in an async action
must also be wrapped in their own `startTransition`, due to the absence
of an AsyncContext mechanism in browsers today.

---

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 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.

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

This is basically how flushSync is implemented, too.

DiffTrain build for commit ee5c194.

Author: acdlite

compiled-rn/facebook-fbsource/xplat/js/RKJSModules/vendor/react/react-test-renderer/cjs/ReactTestRenderer-dev.js

@@ -7,7 +7,7 @@
  * @noflow
  * @nolint
  * @preventMunge
- * @generated SignedSource<<7d577e8126a68c701e64111a9f2bbd2a>>
+ * @generated SignedSource<<8b9b9798442c5728adfabb7cd8ca5f24>>
  */
 
 'use strict';
@@ -8576,9 +8576,7 @@ function runActionStateAction(actionQueue, setPendingState, setState, payload) {
   var prevState = actionQueue.state; // This is a fork of startTransition
 
   var prevTransition = ReactSharedInternals.T;
-  var currentTransition = {
-    _callbacks: new Set()
-  };
+  var currentTransition = {};
   ReactSharedInternals.T = currentTransition;
 
   {
@@ -8591,11 +8589,15 @@ function runActionStateAction(actionQueue, setPendingState, setState, payload) {
 
   try {
     var returnValue = action(prevState, payload);
+    var onStartTransitionFinish = ReactSharedInternals.S;
+
+    if (onStartTransitionFinish !== null) {
+      onStartTransitionFinish(currentTransition, returnValue);
+    }
 
     if (returnValue !== null && typeof returnValue === 'object' && // $FlowFixMe[method-unbinding]
     typeof returnValue.then === 'function') {
-      var thenable = returnValue;
-      notifyTransitionCallbacks(currentTransition, thenable); // Attach a listener to read the return state of the action. As soon as
+      var thenable = returnValue; // 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.
 
       thenable.then(function (nextState) {
@@ -9100,9 +9102,7 @@ function startTransition(fiber, queue, pendingState, finishedState, callback, op
   var previousPriority = getCurrentUpdatePriority();
   setCurrentUpdatePriority(higherEventPriority(previousPriority, ContinuousEventPriority));
   var prevTransition = ReactSharedInternals.T;
-  var currentTransition = {
-    _callbacks: new Set()
-  };
+  var currentTransition = {};
 
   {
     // We don't really need to use an optimistic update here, because we
@@ -9121,7 +9121,12 @@ function startTransition(fiber, queue, pendingState, finishedState, callback, op
 
   try {
     if (enableAsyncActions) {
-      var returnValue = callback(); // Check if we're inside an async action scope. If so, we'll entangle
+      var returnValue = callback();
+      var 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.
       //
       // If we're not already inside an async action scope, and this action is
@@ -9130,9 +9135,9 @@ function startTransition(fiber, queue, pendingState, finishedState, callback, op
       // In the async case, the resulting render will suspend until the async
       // action scope has finished.
 
+
       if (returnValue !== null && typeof returnValue === 'object' && typeof returnValue.then === 'function') {
-        var thenable = returnValue;
-        notifyTransitionCallbacks(currentTransition, thenable); // Create a thenable that resolves to `finishedState` once the async
+        var thenable = returnValue; // Create a thenable that resolves to `finishedState` once the async
         // action has completed.
 
         var thenableForFinishedState = chainThenableValue(thenable, finishedState);
@@ -15141,30 +15146,42 @@ function popCacheProvider(workInProgress, cache) {
   popProvider(CacheContext, workInProgress);
 }
 
-function requestCurrentTransition() {
-  var 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);
+// 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.
+
+var prevOnStartTransitionFinish = ReactSharedInternals.S;
+
+ReactSharedInternals.S = function onStartTransitionFinishForReconciler(transition, returnValue) {
+  if (typeof returnValue === 'object' && returnValue !== null && typeof returnValue.then === 'function') {
+    // This is an async action
+    var thenable = returnValue;
+    entangleAsyncAction(transition, thenable);
   }
 
-  return transition;
-}
-
-function handleAsyncAction(transition, thenable) {
-  {
-    // This is an async action.
-    entangleAsyncAction(transition, thenable);
+  if (prevOnStartTransitionFinish !== null) {
+    prevOnStartTransitionFinish(transition, returnValue);
   }
-}
+};
 
-function notifyTransitionCallbacks(transition, returnValue) {
-  var callbacks = transition._callbacks;
-  callbacks.forEach(function (callback) {
-    return callback(transition, returnValue);
-  });
+function requestCurrentTransition() {
+  return ReactSharedInternals.T;
 } // When retrying a Suspense/Offscreen boundary, we restore the cache that was
 // used during the previous render by placing it here, on the stack.
 
@@ -23446,7 +23463,7 @@ identifierPrefix, onUncaughtError, onCaughtError, onRecoverableError, transition
   return root;
 }
 
-var ReactVersion = '19.0.0-rc-c940c483';
+var ReactVersion = '19.0.0-rc-5a18a922';
 
 /*
  * The `'' + value` pattern (used in perf-sensitive code) throws for Symbol

compiled-rn/facebook-fbsource/xplat/js/RKJSModules/vendor/react/react-test-renderer/cjs/ReactTestRenderer-prod.js

@@ -7,7 +7,7 @@
  * @noflow
  * @nolint
  * @preventMunge
- * @generated SignedSource<<8fc32585cc376c14df2f02921b08d02e>>
+ * @generated SignedSource<<e0bd67bfbe6ced8b6e410eb3974b80d7>>
  */
 
 "use strict";
@@ -2780,16 +2780,18 @@ function runActionStateAction(actionQueue, setPendingState, setState, payload) {
   var action = actionQueue.action,
     prevState = actionQueue.state,
     prevTransition = ReactSharedInternals.T,
-    currentTransition = { _callbacks: new Set() };
+    currentTransition = {};
   ReactSharedInternals.T = currentTransition;
   setPendingState(!0);
   try {
-    var returnValue = action(prevState, payload);
+    var returnValue = action(prevState, payload),
+      onStartTransitionFinish = ReactSharedInternals.S;
+    null !== onStartTransitionFinish &&
+      onStartTransitionFinish(currentTransition, returnValue);
     null !== returnValue &&
     "object" === typeof returnValue &&
     "function" === typeof returnValue.then
-      ? (notifyTransitionCallbacks(currentTransition, returnValue),
-        returnValue.then(
+      ? (returnValue.then(
           function (nextState) {
             actionQueue.state = nextState;
             finishRunningActionStateAction(
@@ -3049,17 +3051,19 @@ function startTransition(fiber, queue, pendingState, finishedState, callback) {
   currentUpdatePriority =
     0 !== previousPriority && 8 > previousPriority ? previousPriority : 8;
   var prevTransition = ReactSharedInternals.T,
-    currentTransition = { _callbacks: new Set() };
+    currentTransition = {};
   ReactSharedInternals.T = currentTransition;
   dispatchOptimisticSetState(fiber, !1, queue, pendingState);
   try {
-    var returnValue = callback();
+    var returnValue = callback(),
+      onStartTransitionFinish = ReactSharedInternals.S;
+    null !== onStartTransitionFinish &&
+      onStartTransitionFinish(currentTransition, returnValue);
     if (
       null !== returnValue &&
       "object" === typeof returnValue &&
       "function" === typeof returnValue.then
     ) {
-      notifyTransitionCallbacks(currentTransition, returnValue);
       var thenableForFinishedState = chainThenableValue(
         returnValue,
         finishedState
@@ -3160,7 +3164,6 @@ function dispatchSetState(fiber, queue, action) {
   }
 }
 function dispatchOptimisticSetState(fiber, throwIfDuringRender, queue, action) {
-  requestCurrentTransition();
   action = {
     lane: 2,
     revertLane: requestTransitionLane(),
@@ -5602,19 +5605,15 @@ function releaseCache(cache) {
       cache.controller.abort();
     });
 }
-function requestCurrentTransition() {
-  var transition = ReactSharedInternals.T;
-  null !== transition && transition._callbacks.add(handleAsyncAction);
-  return transition;
-}
-function handleAsyncAction(transition, thenable) {
-  entangleAsyncAction(transition, thenable);
-}
-function notifyTransitionCallbacks(transition, returnValue) {
-  transition._callbacks.forEach(function (callback) {
-    return callback(transition, returnValue);
-  });
-}
+var prevOnStartTransitionFinish = ReactSharedInternals.S;
+ReactSharedInternals.S = function (transition, returnValue) {
+  "object" === typeof returnValue &&
+    null !== returnValue &&
+    "function" === typeof returnValue.then &&
+    entangleAsyncAction(transition, returnValue);
+  null !== prevOnStartTransitionFinish &&
+    prevOnStartTransitionFinish(transition, returnValue);
+};
 var resumedCache = createCursor(null);
 function peekCacheFromPool() {
   var cacheResumedFromPreviousRender = resumedCache.current;
@@ -7573,7 +7572,7 @@ function requestUpdateLane(fiber) {
   if (0 === (fiber.mode & 1)) return 2;
   if (0 !== (executionContext & 2) && 0 !== workInProgressRootRenderLanes)
     return workInProgressRootRenderLanes & -workInProgressRootRenderLanes;
-  if (null !== requestCurrentTransition())
+  if (null !== ReactSharedInternals.T)
     return (
       (fiber = currentEntangledLane),
       0 !== fiber ? fiber : requestTransitionLane()
@@ -9301,7 +9300,7 @@ var devToolsConfig$jscomp$inline_1042 = {
     throw Error("TestRenderer does not support findFiberByHostInstance()");
   },
   bundleType: 0,
-  version: "19.0.0-rc-2bec340c",
+  version: "19.0.0-rc-9c2862e7",
   rendererPackageName: "react-test-renderer"
 };
 var internals$jscomp$inline_1229 = {
@@ -9332,7 +9331,7 @@ var internals$jscomp$inline_1229 = {
   scheduleRoot: null,
   setRefreshHandler: null,
   getCurrentFiber: null,
-  reconcilerVersion: "19.0.0-rc-2bec340c"
+  reconcilerVersion: "19.0.0-rc-9c2862e7"
 };
 if ("undefined" !== typeof __REACT_DEVTOOLS_GLOBAL_HOOK__) {
   var hook$jscomp$inline_1230 = __REACT_DEVTOOLS_GLOBAL_HOOK__;

compiled-rn/facebook-fbsource/xplat/js/RKJSModules/vendor/react/react-test-renderer/cjs/ReactTestRenderer-profiling.js

@@ -7,7 +7,7 @@
  * @noflow
  * @nolint
  * @preventMunge
- * @generated SignedSource<<f60fad93c112dcc72285bcbbda070adc>>
+ * @generated SignedSource<<cd5cd2d51aad2dcc6d485f9a9bc90b90>>
  */
 
 "use strict";
@@ -2868,16 +2868,18 @@ function runActionStateAction(actionQueue, setPendingState, setState, payload) {
   var action = actionQueue.action,
     prevState = actionQueue.state,
     prevTransition = ReactSharedInternals.T,
-    currentTransition = { _callbacks: new Set() };
+    currentTransition = {};
   ReactSharedInternals.T = currentTransition;
   setPendingState(!0);
   try {
-    var returnValue = action(prevState, payload);
+    var returnValue = action(prevState, payload),
+      onStartTransitionFinish = ReactSharedInternals.S;
+    null !== onStartTransitionFinish &&
+      onStartTransitionFinish(currentTransition, returnValue);
     null !== returnValue &&
     "object" === typeof returnValue &&
     "function" === typeof returnValue.then
-      ? (notifyTransitionCallbacks(currentTransition, returnValue),
-        returnValue.then(
+      ? (returnValue.then(
           function (nextState) {
             actionQueue.state = nextState;
             finishRunningActionStateAction(
@@ -3139,17 +3141,19 @@ function startTransition(fiber, queue, pendingState, finishedState, callback) {
   currentUpdatePriority =
     0 !== previousPriority && 8 > previousPriority ? previousPriority : 8;
   var prevTransition = ReactSharedInternals.T,
-    currentTransition = { _callbacks: new Set() };
+    currentTransition = {};
   ReactSharedInternals.T = currentTransition;
   dispatchOptimisticSetState(fiber, !1, queue, pendingState);
   try {
-    var returnValue = callback();
+    var returnValue = callback(),
+      onStartTransitionFinish = ReactSharedInternals.S;
+    null !== onStartTransitionFinish &&
+      onStartTransitionFinish(currentTransition, returnValue);
     if (
       null !== returnValue &&
       "object" === typeof returnValue &&
       "function" === typeof returnValue.then
     ) {
-      notifyTransitionCallbacks(currentTransition, returnValue);
       var thenableForFinishedState = chainThenableValue(
         returnValue,
         finishedState
@@ -3252,7 +3256,6 @@ function dispatchSetState(fiber, queue, action) {
   markStateUpdateScheduled(fiber, lane);
 }
 function dispatchOptimisticSetState(fiber, throwIfDuringRender, queue, action) {
-  requestCurrentTransition();
   action = {
     lane: 2,
     revertLane: requestTransitionLane(),
@@ -5809,19 +5812,15 @@ function releaseCache(cache) {
       cache.controller.abort();
     });
 }
-function requestCurrentTransition() {
-  var transition = ReactSharedInternals.T;
-  null !== transition && transition._callbacks.add(handleAsyncAction);
-  return transition;
-}
-function handleAsyncAction(transition, thenable) {
-  entangleAsyncAction(transition, thenable);
-}
-function notifyTransitionCallbacks(transition, returnValue) {
-  transition._callbacks.forEach(function (callback) {
-    return callback(transition, returnValue);
-  });
-}
+var prevOnStartTransitionFinish = ReactSharedInternals.S;
+ReactSharedInternals.S = function (transition, returnValue) {
+  "object" === typeof returnValue &&
+    null !== returnValue &&
+    "function" === typeof returnValue.then &&
+    entangleAsyncAction(transition, returnValue);
+  null !== prevOnStartTransitionFinish &&
+    prevOnStartTransitionFinish(transition, returnValue);
+};
 var resumedCache = createCursor(null);
 function peekCacheFromPool() {
   var cacheResumedFromPreviousRender = resumedCache.current;
@@ -8056,7 +8055,7 @@ function requestUpdateLane(fiber) {
   if (0 === (fiber.mode & 1)) return 2;
   if (0 !== (executionContext & 2) && 0 !== workInProgressRootRenderLanes)
     return workInProgressRootRenderLanes & -workInProgressRootRenderLanes;
-  if (null !== requestCurrentTransition())
+  if (null !== ReactSharedInternals.T)
     return (
       (fiber = currentEntangledLane),
       0 !== fiber ? fiber : requestTransitionLane()
@@ -9944,7 +9943,7 @@ var devToolsConfig$jscomp$inline_1105 = {
     throw Error("TestRenderer does not support findFiberByHostInstance()");
   },
   bundleType: 0,
-  version: "19.0.0-rc-e8e0cdf6",
+  version: "19.0.0-rc-6cf53e9c",
   rendererPackageName: "react-test-renderer"
 };
 (function (internals) {
@@ -9988,7 +9987,7 @@ var devToolsConfig$jscomp$inline_1105 = {
   scheduleRoot: null,
   setRefreshHandler: null,
   getCurrentFiber: null,
-  reconcilerVersion: "19.0.0-rc-e8e0cdf6"
+  reconcilerVersion: "19.0.0-rc-6cf53e9c"
 });
 exports._Scheduler = Scheduler;
 exports.act = act;