Async action support for React.startTransition (#28097)

This adds support for async actions to the "isomorphic" version of
startTransition (i.e. the one exported by the "react" package).
Previously, async actions were only supported by the startTransition
that is returned from the useTransition hook.

The interesting part about the isomorphic startTransition is that it's
not associated with any particular root. It must work with updates to
arbitrary roots, or even arbitrary React renderers in the same app. (For
example, both React DOM and React Three Fiber.)

The idea is that React.startTransition should behave as if every root
had an implicit useTransition hook, and you composed together all the
startTransitions provided by those hooks. Multiple updates to the same
root will be batched together. However, updates to one root will not be
batched with updates to other roots.

Features like useOptimistic work the same as with the hook version.

There is one difference from from the hook version of startTransition:
an error triggered inside an async action cannot be captured by an error
boundary, because it's not associated with any particular part of the
tree. You should handle errors the same way you would in a regular
event, e.g. with a global error event handler, or with a local
`try/catch`.

DiffTrain build for commit 85b296e.

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

@@ -7,7 +7,7 @@
  * @noflow
  * @nolint
  * @preventMunge
- * @generated SignedSource<<6f91cae5b10795257ddec042012ed6d4>>
+ * @generated SignedSource<<44707665e2d49df89524e6821112ab6c>>
  */
 
 "use strict";
@@ -3469,7 +3469,11 @@ if (__DEV__) {
       }
     }
 
-    function requestTransitionLane() {
+    function requestTransitionLane( // This argument isn't used, it's only here to encourage the caller to
+      // check that it's inside a transition before calling this function.
+      // TODO: Make this non-nullable. Requires a tweak to useOptimistic.
+      transition
+    ) {
       // The algorithm for assigning an update to a lane should be stable for all
       // updates at the same priority within the same event. To do this, the
       // inputs to the algorithm must be the same.
@@ -3502,7 +3506,7 @@ if (__DEV__) {
     // until the async action scope has completed.
 
     var currentEntangledActionThenable = null;
-    function entangleAsyncAction(thenable) {
+    function entangleAsyncAction(transition, thenable) {
       // `thenable` is the return value of the async action scope function. Create
       // a combined thenable that resolves once every entangled scope function
       // has finished.
@@ -7738,16 +7742,7 @@ if (__DEV__) {
             markSkippedUpdateLanes(updateLane);
           } else {
             // This update does have sufficient priority.
-            // Check if this update is part of a pending async action. If so,
-            // we'll need to suspend until the action has finished, so that it's
-            // batched together with future updates in the same action.
-            if (
-              updateLane !== NoLane &&
-              updateLane === peekEntangledActionLane()
-            ) {
-              didReadFromEntangledAsyncAction = true;
-            } // Check if this is an optimistic update.
-
+            // Check if this is an optimistic update.
             var revertLane = update.revertLane;
 
             if (revertLane === NoLane) {
@@ -7767,6 +7762,12 @@ if (__DEV__) {
                   next: null
                 };
                 newBaseQueueLast = newBaseQueueLast.next = _clone;
+              } // Check if this update is part of a pending async action. If so,
+              // we'll need to suspend until the action has finished, so that it's
+              // batched together with future updates in the same action.
+
+              if (updateLane === peekEntangledActionLane()) {
+                didReadFromEntangledAsyncAction = true;
               }
             } else {
               // This is an optimistic update. If the "revert" priority is
@@ -7777,7 +7778,14 @@ if (__DEV__) {
                 // The transition that this optimistic update is associated with
                 // has finished. Pretend the update doesn't exist by skipping
                 // over it.
-                update = update.next;
+                update = update.next; // Check if this update is part of a pending async action. If so,
+                // we'll need to suspend until the action has finished, so that it's
+                // batched together with future updates in the same action.
+
+                if (revertLane === peekEntangledActionLane()) {
+                  didReadFromEntangledAsyncAction = true;
+                }
+
                 continue;
               } else {
                 var _clone2 = {
@@ -8288,15 +8296,18 @@ if (__DEV__) {
       var prevState = actionQueue.state; // This is a fork of startTransition
 
       var prevTransition = ReactCurrentBatchConfig$2.transition;
-      ReactCurrentBatchConfig$2.transition = {};
-      var currentTransition = ReactCurrentBatchConfig$2.transition;
+      var currentTransition = {
+        _callbacks: new Set()
+      };
+      ReactCurrentBatchConfig$2.transition = currentTransition;
 
       {
         ReactCurrentBatchConfig$2.transition._updatedFibers = new Set();
       }
 
       try {
         var returnValue = action(prevState, payload);
+        notifyTransitionCallbacks(currentTransition, returnValue);
 
         if (
           returnValue !== null &&
@@ -8315,7 +8326,6 @@ if (__DEV__) {
               return finishRunningFormStateAction(actionQueue, setState);
             }
           );
-          entangleAsyncAction(thenable);
           setState(thenable);
         } else {
           setState(returnValue);
@@ -8876,7 +8886,9 @@ if (__DEV__) {
         higherEventPriority(previousPriority, ContinuousEventPriority)
       );
       var prevTransition = ReactCurrentBatchConfig$2.transition;
-      var currentTransition = {};
+      var currentTransition = {
+        _callbacks: new Set()
+      };
 
       {
         // We don't really need to use an optimistic update here, because we
@@ -8895,7 +8907,8 @@ if (__DEV__) {
 
       try {
         if (enableAsyncActions) {
-          var returnValue = callback(); // Check if we're inside an async action scope. If so, we'll entangle
+          var returnValue = callback();
+          notifyTransitionCallbacks(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
@@ -8909,8 +8922,7 @@ if (__DEV__) {
             typeof returnValue === "object" &&
             typeof returnValue.then === "function"
           ) {
-            var thenable = returnValue;
-            entangleAsyncAction(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(
@@ -9214,8 +9226,10 @@ if (__DEV__) {
       queue,
       action
     ) {
+      var transition = requestCurrentTransition();
+
       {
-        if (ReactCurrentBatchConfig$2.transition === null) {
+        if (transition === null) {
           // An optimistic update occurred, but startTransition is not on the stack.
           // There are two likely scenarios.
           // One possibility is that the optimistic update is triggered by a regular
@@ -16268,9 +16282,35 @@ if (__DEV__) {
 
     var ReactCurrentBatchConfig$1 =
       ReactSharedInternals.ReactCurrentBatchConfig;
-    var NoTransition = null;
     function requestCurrentTransition() {
-      return ReactCurrentBatchConfig$1.transition;
+      var transition = ReactCurrentBatchConfig$1.transition;
+
+      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(handleTransitionScopeResult);
+      }
+
+      return transition;
+    }
+
+    function handleTransitionScopeResult(transition, returnValue) {
+      if (
+        returnValue !== null &&
+        typeof returnValue === "object" &&
+        typeof returnValue.then === "function"
+      ) {
+        // This is an async action.
+        var thenable = returnValue;
+        entangleAsyncAction(transition, thenable);
+      }
+    }
+
+    function notifyTransitionCallbacks(transition, returnValue) {
+      var callbacks = transition._callbacks;
+      callbacks.forEach(function (callback) {
+        return callback(transition, returnValue);
+      });
     } // When retrying a Suspense/Offscreen boundary, we restore the cache that was
     // used during the previous render by placing it here, on the stack.
 
@@ -21566,17 +21606,17 @@ if (__DEV__) {
         return pickArbitraryLane(workInProgressRootRenderLanes);
       }
 
-      var isTransition = requestCurrentTransition() !== NoTransition;
+      var transition = requestCurrentTransition();
 
-      if (isTransition) {
-        if (ReactCurrentBatchConfig.transition !== null) {
-          var transition = ReactCurrentBatchConfig.transition;
+      if (transition !== null) {
+        {
+          var batchConfigTransition = ReactCurrentBatchConfig.transition;
 
-          if (!transition._updatedFibers) {
-            transition._updatedFibers = new Set();
+          if (!batchConfigTransition._updatedFibers) {
+            batchConfigTransition._updatedFibers = new Set();
           }
 
-          transition._updatedFibers.add(fiber);
+          batchConfigTransition._updatedFibers.add(fiber);
         }
 
         var actionScopeLane = peekEntangledActionLane();
@@ -21643,7 +21683,7 @@ if (__DEV__) {
           workInProgressDeferredLane = OffscreenLane;
         } else {
           // Everything else is spawned as a transition.
-          workInProgressDeferredLane = requestTransitionLane();
+          workInProgressDeferredLane = claimNextTransitionLane();
         }
       } // Mark the parent Suspense boundary so it knows to spawn the deferred lane.
 
@@ -25572,7 +25612,7 @@ if (__DEV__) {
       return root;
     }
 
-    var ReactVersion = "18.3.0-canary-382190c59-20240125";
+    var ReactVersion = "18.3.0-canary-85b296e9b-20240125";
 
     // Might add PROFILE later.