Async action support for React.startTransition

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

Author: acdlite

packages/react-reconciler/src/ReactFiberAsyncAction.js

@@ -13,6 +13,7 @@ import type {
   RejectedThenable,
 } from 'shared/ReactTypes';
 import type {Lane} from './ReactFiberLane';
+import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';
 
 import {requestTransitionLane} from './ReactFiberRootScheduler';
 import {NoLane} from './ReactFiberLane';
@@ -36,15 +37,18 @@ let currentEntangledLane: Lane = NoLane;
 // until the async action scope has completed.
 let currentEntangledActionThenable: Thenable<void> | null = null;
 
-export function entangleAsyncAction<S>(thenable: Thenable<S>): Thenable<S> {
+export function entangleAsyncAction<S>(
+  transition: BatchConfigTransition,
+  thenable: Thenable<S>,
+): Thenable<S> {
   // `thenable` is the return value of the async action scope function. Create
   // a combined thenable that resolves once every entangled scope function
   // has finished.
   if (currentEntangledListeners === null) {
     // There's no outer async action scope. Create a new one.
     const entangledListeners = (currentEntangledListeners = []);
     currentEntangledPendingCount = 0;
-    currentEntangledLane = requestTransitionLane();
+    currentEntangledLane = requestTransitionLane(transition);
     const entangledThenable: Thenable<void> = {
       status: 'pending',
       value: undefined,

packages/react-reconciler/src/ReactFiberHooks.js

@@ -145,14 +145,17 @@ import {
 import type {ThenableState} from './ReactFiberThenable';
 import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';
 import {
-  entangleAsyncAction,
   peekEntangledActionLane,
   peekEntangledActionThenable,
   chainThenableValue,
 } from './ReactFiberAsyncAction';
 import {HostTransitionContext} from './ReactFiberHostContext';
 import {requestTransitionLane} from './ReactFiberRootScheduler';
 import {isCurrentTreeHidden} from './ReactFiberHiddenContext';
+import {
+  notifyTransitionCallbacks,
+  requestCurrentTransition,
+} from './ReactFiberTransition';
 
 const {ReactCurrentDispatcher, ReactCurrentBatchConfig} = ReactSharedInternals;
 
@@ -1319,13 +1322,6 @@ function updateReducerImpl<S, A>(
       } 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.
         const revertLane = update.revertLane;
         if (!enableAsyncActions || revertLane === NoLane) {
@@ -1346,6 +1342,13 @@ function updateReducerImpl<S, A>(
             };
             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
           // sufficient, don't apply the update. Otherwise, apply the update,
@@ -1356,6 +1359,13 @@ function updateReducerImpl<S, A>(
             // has finished. Pretend the update doesn't exist by skipping
             // over it.
             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 {
             const clone: Update<S, A> = {
@@ -1964,13 +1974,17 @@ function runFormStateAction<S, P>(
 
   // This is a fork of startTransition
   const prevTransition = ReactCurrentBatchConfig.transition;
-  ReactCurrentBatchConfig.transition = ({}: BatchConfigTransition);
-  const currentTransition = ReactCurrentBatchConfig.transition;
+  const currentTransition: BatchConfigTransition = {
+    _callbacks: new Set<(BatchConfigTransition, mixed) => mixed>(),
+  };
+  ReactCurrentBatchConfig.transition = currentTransition;
   if (__DEV__) {
     ReactCurrentBatchConfig.transition._updatedFibers = new Set();
   }
   try {
     const returnValue = action(prevState, payload);
+    notifyTransitionCallbacks(currentTransition, returnValue);
+
     if (
       returnValue !== null &&
       typeof returnValue === 'object' &&
@@ -1989,7 +2003,6 @@ function runFormStateAction<S, P>(
         () => finishRunningFormStateAction(actionQueue, (setState: any)),
       );
 
-      entangleAsyncAction<Awaited<S>>(thenable);
       setState((thenable: any));
     } else {
       setState((returnValue: any));
@@ -2808,7 +2821,9 @@ function startTransition<S>(
   );
 
   const prevTransition = ReactCurrentBatchConfig.transition;
-  const currentTransition: BatchConfigTransition = {};
+  const currentTransition: BatchConfigTransition = {
+    _callbacks: new Set<(BatchConfigTransition, mixed) => mixed>(),
+  };
 
   if (enableAsyncActions) {
     // We don't really need to use an optimistic update here, because we
@@ -2839,6 +2854,7 @@ function startTransition<S>(
   try {
     if (enableAsyncActions) {
       const 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.
@@ -2854,7 +2870,6 @@ function startTransition<S>(
         typeof returnValue.then === 'function'
       ) {
         const thenable = ((returnValue: any): Thenable<mixed>);
-        entangleAsyncAction<mixed>(thenable);
         // Create a thenable that resolves to `finishedState` once the async
         // action has completed.
         const thenableForFinishedState = chainThenableValue(
@@ -3281,8 +3296,10 @@ function dispatchOptimisticSetState<S, A>(
   queue: UpdateQueue<S, A>,
   action: A,
 ): void {
+  const transition = requestCurrentTransition();
+
   if (__DEV__) {
-    if (ReactCurrentBatchConfig.transition === null) {
+    if (transition === null) {
       // An optimistic update occurred, but startTransition is not on the stack.
       // There are two likely scenarios.
 
@@ -3323,7 +3340,7 @@ function dispatchOptimisticSetState<S, A>(
     lane: SyncLane,
     // After committing, the optimistic update is "reverted" using the same
     // lane as the transition it's associated with.
-    revertLane: requestTransitionLane(),
+    revertLane: requestTransitionLane(transition),
     action,
     hasEagerState: false,
     eagerState: null,

packages/react-reconciler/src/ReactFiberRootScheduler.js

@@ -10,6 +10,7 @@
 import type {FiberRoot} from './ReactInternalTypes';
 import type {Lane} from './ReactFiberLane';
 import type {PriorityLevel} from 'scheduler/src/SchedulerPriorities';
+import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';
 
 import {enableDeferRootSchedulingToMicrotask} from 'shared/ReactFeatureFlags';
 import {
@@ -492,7 +493,12 @@ function scheduleImmediateTask(cb: () => mixed) {
   }
 }
 
-export function requestTransitionLane(): Lane {
+export 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: BatchConfigTransition | null,
+): Lane {
   // 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.

packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js

@@ -36,6 +36,7 @@ export type PendingTransitionCallbacks = {
   markerComplete: Map<string, Set<Transition>> | null,
 };
 
+// TODO: Unclear to me why these are separate types
 export type Transition = {
   name: string,
   startTime: number,
@@ -45,6 +46,7 @@ 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

@@ -7,12 +7,20 @@
  * @flow
  */
 import type {Fiber, FiberRoot} from './ReactInternalTypes';
+import type {Thenable} from 'shared/ReactTypes';
 import type {Lanes} from './ReactFiberLane';
 import type {StackCursor} from './ReactFiberStack';
 import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent';
-import type {Transition} from './ReactFiberTracingMarkerComponent';
+import type {
+  BatchConfigTransition,
+  Transition,
+} from './ReactFiberTracingMarkerComponent';
 
-import {enableCache, enableTransitionTracing} from 'shared/ReactFeatureFlags';
+import {
+  enableCache,
+  enableTransitionTracing,
+  enableAsyncActions,
+} from 'shared/ReactFeatureFlags';
 import {isPrimaryRenderer} from './ReactFiberConfig';
 import {createCursor, push, pop} from './ReactFiberStack';
 import {
@@ -26,13 +34,44 @@ import {
 } from './ReactFiberCacheComponent';
 
 import ReactSharedInternals from 'shared/ReactSharedInternals';
+import {entangleAsyncAction} from './ReactFiberAsyncAction';
 
 const {ReactCurrentBatchConfig} = ReactSharedInternals;
 
 export const NoTransition = null;
 
-export function requestCurrentTransition(): Transition | null {
-  return ReactCurrentBatchConfig.transition;
+export function requestCurrentTransition(): BatchConfigTransition | null {
+  const transition = ReactCurrentBatchConfig.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: BatchConfigTransition,
+  returnValue: mixed,
+): void {
+  if (
+    enableAsyncActions &&
+    returnValue !== null &&
+    typeof returnValue === 'object' &&
+    typeof returnValue.then === 'function'
+  ) {
+    // This is an async action.
+    const thenable: Thenable<mixed> = (returnValue: any);
+    entangleAsyncAction(transition, thenable);
+  }
+}
+
+export function notifyTransitionCallbacks(
+  transition: BatchConfigTransition,
+  returnValue: mixed,
+) {
+  const callbacks = transition._callbacks;
+  callbacks.forEach(callback => callback(transition, returnValue));
 }
 
 // When retrying a Suspense/Offscreen boundary, we restore the cache that was

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -161,6 +161,7 @@ import {
   OffscreenLane,
   SyncUpdateLanes,
   UpdateLanes,
+  claimNextTransitionLane,
 } from './ReactFiberLane';
 import {
   DiscreteEventPriority,
@@ -170,7 +171,7 @@ import {
   lowerEventPriority,
   lanesToEventPriority,
 } from './ReactEventPriorities';
-import {requestCurrentTransition, NoTransition} from './ReactFiberTransition';
+import {requestCurrentTransition} from './ReactFiberTransition';
 import {
   SelectiveHydrationException,
   beginWork as originalBeginWork,
@@ -633,15 +634,15 @@ export function requestUpdateLane(fiber: Fiber): Lane {
     return pickArbitraryLane(workInProgressRootRenderLanes);
   }
 
-  const isTransition = requestCurrentTransition() !== NoTransition;
-  if (isTransition) {
-    if (__DEV__ && ReactCurrentBatchConfig.transition !== null) {
-      const transition = ReactCurrentBatchConfig.transition;
-      if (!transition._updatedFibers) {
-        transition._updatedFibers = new Set();
+  const transition = requestCurrentTransition();
+  if (transition !== null) {
+    if (__DEV__) {
+      const batchConfigTransition = ReactCurrentBatchConfig.transition;
+      if (!batchConfigTransition._updatedFibers) {
+        batchConfigTransition._updatedFibers = new Set();
       }
 
-      transition._updatedFibers.add(fiber);
+      batchConfigTransition._updatedFibers.add(fiber);
     }
 
     const actionScopeLane = peekEntangledActionLane();
@@ -651,7 +652,7 @@ export function requestUpdateLane(fiber: Fiber): Lane {
       : // We may or may not be inside an async action scope. If we are, this
         // is the first update in that scope. Either way, we need to get a
         // fresh transition lane.
-        requestTransitionLane();
+        requestTransitionLane(transition);
   }
 
   // Updates originating inside certain React methods, like flushSync, have
@@ -712,7 +713,7 @@ export function requestDeferredLane(): Lane {
       workInProgressDeferredLane = OffscreenLane;
     } else {
       // Everything else is spawned as a transition.
-      workInProgressDeferredLane = requestTransitionLane();
+      workInProgressDeferredLane = claimNextTransitionLane();
     }
   }