From 2a83d9a511b0a9be0d62193c36cfd9e424bd58df Mon Sep 17 00:00:00 2001 From: Andrew Clark Date: Thu, 25 Jan 2024 13:26:23 -0500 Subject: [PATCH] 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`. --- .../src/ReactFiberAsyncAction.js | 8 +- .../react-reconciler/src/ReactFiberHooks.js | 47 ++++++---- .../src/ReactFiberRootScheduler.js | 8 +- .../src/ReactFiberTracingMarkerComponent.js | 2 + .../src/ReactFiberTransition.js | 47 +++++++++- .../src/ReactFiberWorkLoop.js | 21 ++--- .../src/__tests__/ReactAsyncActions-test.js | 85 +++++++++++++++++++ packages/react/src/ReactStartTransition.js | 11 ++- 8 files changed, 195 insertions(+), 34 deletions(-) diff --git a/packages/react-reconciler/src/ReactFiberAsyncAction.js b/packages/react-reconciler/src/ReactFiberAsyncAction.js index 98f7ca175c0a8..ec53c7d80346c 100644 --- a/packages/react-reconciler/src/ReactFiberAsyncAction.js +++ b/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,7 +37,10 @@ let currentEntangledLane: Lane = NoLane; // until the async action scope has completed. let currentEntangledActionThenable: Thenable | null = null; -export function entangleAsyncAction(thenable: Thenable): Thenable { +export function entangleAsyncAction( + transition: BatchConfigTransition, + thenable: Thenable, +): 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. @@ -44,7 +48,7 @@ export function entangleAsyncAction(thenable: Thenable): Thenable { // There's no outer async action scope. Create a new one. const entangledListeners = (currentEntangledListeners = []); currentEntangledPendingCount = 0; - currentEntangledLane = requestTransitionLane(); + currentEntangledLane = requestTransitionLane(transition); const entangledThenable: Thenable = { status: 'pending', value: undefined, diff --git a/packages/react-reconciler/src/ReactFiberHooks.js b/packages/react-reconciler/src/ReactFiberHooks.js index 42a9eaa4f477d..dd3cf8c273d3a 100644 --- a/packages/react-reconciler/src/ReactFiberHooks.js +++ b/packages/react-reconciler/src/ReactFiberHooks.js @@ -145,7 +145,6 @@ import { import type {ThenableState} from './ReactFiberThenable'; import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent'; import { - entangleAsyncAction, peekEntangledActionLane, peekEntangledActionThenable, chainThenableValue, @@ -153,6 +152,10 @@ import { 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( } 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( }; 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( // 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 = { @@ -1964,13 +1974,17 @@ function runFormStateAction( // 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( () => finishRunningFormStateAction(actionQueue, (setState: any)), ); - entangleAsyncAction>(thenable); setState((thenable: any)); } else { setState((returnValue: any)); @@ -2808,7 +2821,9 @@ function startTransition( ); 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( 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( typeof returnValue.then === 'function' ) { const thenable = ((returnValue: any): Thenable); - entangleAsyncAction(thenable); // Create a thenable that resolves to `finishedState` once the async // action has completed. const thenableForFinishedState = chainThenableValue( @@ -3281,8 +3296,10 @@ function dispatchOptimisticSetState( queue: UpdateQueue, 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( 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, diff --git a/packages/react-reconciler/src/ReactFiberRootScheduler.js b/packages/react-reconciler/src/ReactFiberRootScheduler.js index 19294e7d4880d..da45a9278006e 100644 --- a/packages/react-reconciler/src/ReactFiberRootScheduler.js +++ b/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. diff --git a/packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js b/packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js index 389325f465163..5b6548fadb682 100644 --- a/packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js +++ b/packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js @@ -36,6 +36,7 @@ export type PendingTransitionCallbacks = { markerComplete: Map> | 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, + _callbacks: Set<(BatchConfigTransition, mixed) => mixed>, }; // TODO: Is there a way to not include the tag or name here? diff --git a/packages/react-reconciler/src/ReactFiberTransition.js b/packages/react-reconciler/src/ReactFiberTransition.js index a148375a6708e..d420f67dcfe76 100644 --- a/packages/react-reconciler/src/ReactFiberTransition.js +++ b/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 = (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 diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.js b/packages/react-reconciler/src/ReactFiberWorkLoop.js index d487fdb205c5d..597d0089941a2 100644 --- a/packages/react-reconciler/src/ReactFiberWorkLoop.js +++ b/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(); } } diff --git a/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js b/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js index 1f45fd84d0430..d5003fc49b066 100644 --- a/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js +++ b/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js @@ -1641,4 +1641,89 @@ describe('ReactAsyncActions', () => { expect(root).toMatchRenderedOutput(D); }, ); + + // @gate enableAsyncActions + test('React.startTransition supports async actions', async () => { + const startTransition = React.startTransition; + + function App({text}) { + return ; + } + + const root = ReactNoop.createRoot(); + await act(() => { + root.render(); + }); + assertLog(['A']); + + await act(() => { + startTransition(async () => { + // Update to B + root.render(); + + // There's an async gap before C is updated + await getText('Wait before updating to C'); + root.render(); + + Scheduler.log('Async action ended'); + }); + }); + // The update to B is blocked because the async action hasn't completed yet. + assertLog([]); + expect(root).toMatchRenderedOutput('A'); + + // Finish the async action + await act(() => resolveText('Wait before updating to C')); + + // Now both B and C can finish in a single batch. + assertLog(['Async action ended', 'C']); + expect(root).toMatchRenderedOutput('C'); + }); + + // @gate enableAsyncActions + test('useOptimistic works with async actions passed to React.startTransition', async () => { + const startTransition = React.startTransition; + + let setOptimisticText; + function App({text: canonicalText}) { + const [text, _setOptimisticText] = useOptimistic( + canonicalText, + (_, optimisticText) => `${optimisticText} (loading...)`, + ); + setOptimisticText = _setOptimisticText; + return ( + + + + ); + } + + const root = ReactNoop.createRoot(); + await act(() => { + root.render(); + }); + assertLog(['Initial']); + expect(root).toMatchRenderedOutput(Initial); + + // 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'); + setOptimisticText('Updated'); + await getText('Yield before updating'); + Scheduler.log('Async action ended'); + startTransition(() => root.render()); + }); + }); + // Because the action hasn't finished yet, the optimistic UI is shown. + assertLog(['Async action started', 'Updated (loading...)']); + expect(root).toMatchRenderedOutput(Updated (loading...)); + + // 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(Updated); + }); }); diff --git a/packages/react/src/ReactStartTransition.js b/packages/react/src/ReactStartTransition.js index 490caf5d319e9..ab956a2d2cb38 100644 --- a/packages/react/src/ReactStartTransition.js +++ b/packages/react/src/ReactStartTransition.js @@ -17,7 +17,13 @@ export function startTransition( options?: StartTransitionOptions, ) { const prevTransition = ReactCurrentBatchConfig.transition; - ReactCurrentBatchConfig.transition = ({}: BatchConfigTransition); + // Each renderer registers a callback to receive the return value of + // the scope function. This is used to implement async actions. + const callbacks = new Set<(BatchConfigTransition, mixed) => mixed>(); + const transition: BatchConfigTransition = { + _callbacks: callbacks, + }; + ReactCurrentBatchConfig.transition = transition; const currentTransition = ReactCurrentBatchConfig.transition; if (__DEV__) { @@ -34,7 +40,8 @@ export function startTransition( } try { - scope(); + const returnValue = scope(); + callbacks.forEach(callback => callback(currentTransition, returnValue)); } finally { ReactCurrentBatchConfig.transition = prevTransition;