Async action support for React.startTransition

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();
     }
   }