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;