@samitouri / QOS-React-2 / commits / 85b296e9b6

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

Andrew Clark committed Jan 25, 2024 at 21:54 UTC 85b296e9b6ded4accd9ec3389297f95091fb1ac0

8 files changed +195 -34

packages/react-reconciler/src/ReactFiberAsyncAction.js

+6 -2

@@ -13,6 +13,7 @@ import type {

13 RejectedThenable,

14 } from 'shared/ReactTypes';

15 import type {Lane} from './ReactFiberLane';

16 +import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';

18 import {requestTransitionLane} from './ReactFiberRootScheduler';

19 import {NoLane} from './ReactFiberLane';

@@ -36,7 +37,10 @@ let currentEntangledLane: Lane = NoLane;

37 // until the async action scope has completed.

38 let currentEntangledActionThenable: Thenable<void> | null = null;

39 -export function entangleAsyncAction<S>(thenable: Thenable<S>): Thenable<S> {

40 +export function entangleAsyncAction<S>(

41 + transition: BatchConfigTransition,

42 + thenable: Thenable<S>,

43 +): Thenable<S> {

44 // `thenable` is the return value of the async action scope function. Create

45 // a combined thenable that resolves once every entangled scope function

46 // has finished.

@@ -44,7 +48,7 @@ export function entangleAsyncAction<S>(thenable: Thenable<S>): Thenable<S> {

48 // There's no outer async action scope. Create a new one.

49 const entangledListeners = (currentEntangledListeners = []);

50 currentEntangledPendingCount = 0;

47 - currentEntangledLane = requestTransitionLane();

51 + currentEntangledLane = requestTransitionLane(transition);

52 const entangledThenable: Thenable<void> = {

53 status: 'pending',

54 value: undefined,

packages/react-reconciler/src/ReactFiberHooks.js

+32 -15

@@ -145,7 +145,6 @@ import {

145 import type {ThenableState} from './ReactFiberThenable';

146 import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';

147 import {

148 - entangleAsyncAction,

148 peekEntangledActionLane,

149 peekEntangledActionThenable,

150 chainThenableValue,

@@ -153,6 +152,10 @@ import {

152 import {HostTransitionContext} from './ReactFiberHostContext';

153 import {requestTransitionLane} from './ReactFiberRootScheduler';

154 import {isCurrentTreeHidden} from './ReactFiberHiddenContext';

155 +import {

156 + notifyTransitionCallbacks,

157 + requestCurrentTransition,

158 +} from './ReactFiberTransition';

159

160 const {ReactCurrentDispatcher, ReactCurrentBatchConfig} = ReactSharedInternals;

161

@@ -1319,13 +1322,6 @@ function updateReducerImpl<S, A>(

1322 } else {

1323 // This update does have sufficient priority.

1324

1322 - // Check if this update is part of a pending async action. If so,

1323 - // we'll need to suspend until the action has finished, so that it's

1324 - // batched together with future updates in the same action.

1325 - if (updateLane !== NoLane && updateLane === peekEntangledActionLane()) {

1326 - didReadFromEntangledAsyncAction = true;

1327 - }

1328 -

1325 // Check if this is an optimistic update.

1326 const revertLane = update.revertLane;

1327 if (!enableAsyncActions || revertLane === NoLane) {

@@ -1346,6 +1342,13 @@ function updateReducerImpl<S, A>(

1342 };

1343 newBaseQueueLast = newBaseQueueLast.next = clone;

1344 }

1345 +

1346 + // Check if this update is part of a pending async action. If so,

1347 + // we'll need to suspend until the action has finished, so that it's

1348 + // batched together with future updates in the same action.

1349 + if (updateLane === peekEntangledActionLane()) {

1350 + didReadFromEntangledAsyncAction = true;

1351 + }

1352 } else {

1353 // This is an optimistic update. If the "revert" priority is

1354 // sufficient, don't apply the update. Otherwise, apply the update,

@@ -1356,6 +1359,13 @@ function updateReducerImpl<S, A>(

1359 // has finished. Pretend the update doesn't exist by skipping

1360 // over it.

1361 update = update.next;

1362 +

1363 + // Check if this update is part of a pending async action. If so,

1364 + // we'll need to suspend until the action has finished, so that it's

1365 + // batched together with future updates in the same action.

1366 + if (revertLane === peekEntangledActionLane()) {

1367 + didReadFromEntangledAsyncAction = true;

1368 + }

1369 continue;

1370 } else {

1371 const clone: Update<S, A> = {

@@ -1964,13 +1974,17 @@ function runFormStateAction<S, P>(

1974

1975 // This is a fork of startTransition

1976 const prevTransition = ReactCurrentBatchConfig.transition;

1967 - ReactCurrentBatchConfig.transition = ({}: BatchConfigTransition);

1968 - const currentTransition = ReactCurrentBatchConfig.transition;

1977 + const currentTransition: BatchConfigTransition = {

1978 + _callbacks: new Set<(BatchConfigTransition, mixed) => mixed>(),

1979 + };

1980 + ReactCurrentBatchConfig.transition = currentTransition;

1981 if (__DEV__) {

1982 ReactCurrentBatchConfig.transition._updatedFibers = new Set();

1983 }

1984 try {

1985 const returnValue = action(prevState, payload);

1986 + notifyTransitionCallbacks(currentTransition, returnValue);

1987 +

1988 if (

1989 returnValue !== null &&

1990 typeof returnValue === 'object' &&

@@ -1989,7 +2003,6 @@ function runFormStateAction<S, P>(

2003 () => finishRunningFormStateAction(actionQueue, (setState: any)),

2004 );

2005

1992 - entangleAsyncAction<Awaited<S>>(thenable);

2006 setState((thenable: any));

2007 } else {

2008 setState((returnValue: any));

@@ -2808,7 +2821,9 @@ function startTransition<S>(

2821 );

2822

2823 const prevTransition = ReactCurrentBatchConfig.transition;

2811 - const currentTransition: BatchConfigTransition = {};

2824 + const currentTransition: BatchConfigTransition = {

2825 + _callbacks: new Set<(BatchConfigTransition, mixed) => mixed>(),

2826 + };

2827

2828 if (enableAsyncActions) {

2829 // We don't really need to use an optimistic update here, because we

@@ -2839,6 +2854,7 @@ function startTransition<S>(

2854 try {

2855 if (enableAsyncActions) {

2856 const returnValue = callback();

2857 + notifyTransitionCallbacks(currentTransition, returnValue);

2858

2859 // Check if we're inside an async action scope. If so, we'll entangle

2860 // this new action with the existing scope.

@@ -2854,7 +2870,6 @@ function startTransition<S>(

2870 typeof returnValue.then === 'function'

2871 ) {

2872 const thenable = ((returnValue: any): Thenable<mixed>);

2857 - entangleAsyncAction<mixed>(thenable);

2873 // Create a thenable that resolves to `finishedState` once the async

2874 // action has completed.

2875 const thenableForFinishedState = chainThenableValue(

@@ -3281,8 +3296,10 @@ function dispatchOptimisticSetState<S, A>(

3296 queue: UpdateQueue<S, A>,

3297 action: A,

3298 ): void {

3299 + const transition = requestCurrentTransition();

3300 +

3301 if (__DEV__) {

3285 - if (ReactCurrentBatchConfig.transition === null) {

3302 + if (transition === null) {

3303 // An optimistic update occurred, but startTransition is not on the stack.

3304 // There are two likely scenarios.

3305

@@ -3323,7 +3340,7 @@ function dispatchOptimisticSetState<S, A>(

3340 lane: SyncLane,

3341 // After committing, the optimistic update is "reverted" using the same

3342 // lane as the transition it's associated with.

3326 - revertLane: requestTransitionLane(),

3343 + revertLane: requestTransitionLane(transition),

3344 action,

3345 hasEagerState: false,

3346 eagerState: null,

packages/react-reconciler/src/ReactFiberRootScheduler.js

+7 -1

@@ -10,6 +10,7 @@

10 import type {FiberRoot} from './ReactInternalTypes';

11 import type {Lane} from './ReactFiberLane';

12 import type {PriorityLevel} from 'scheduler/src/SchedulerPriorities';

13 +import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';

15 import {enableDeferRootSchedulingToMicrotask} from 'shared/ReactFeatureFlags';

16 import {

@@ -492,7 +493,12 @@ function scheduleImmediateTask(cb: () => mixed) {

493 }

494 }

495

495 -export function requestTransitionLane(): Lane {

496 +export function requestTransitionLane(

497 + // This argument isn't used, it's only here to encourage the caller to

498 + // check that it's inside a transition before calling this function.

499 + // TODO: Make this non-nullable. Requires a tweak to useOptimistic.

500 + transition: BatchConfigTransition | null,

501 +): Lane {

502 // The algorithm for assigning an update to a lane should be stable for all

503 // updates at the same priority within the same event. To do this, the

504 // inputs to the algorithm must be the same.

packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js

@@ -36,6 +36,7 @@ export type PendingTransitionCallbacks = {

36 markerComplete: Map<string, Set<Transition>> | null,

37 };

39 +// TODO: Unclear to me why these are separate types

40 export type Transition = {

41 name: string,

42 startTime: number,

@@ -45,6 +46,7 @@ export type BatchConfigTransition = {

46 name?: string,

47 startTime?: number,

48 _updatedFibers?: Set<Fiber>,

49 + _callbacks: Set<(BatchConfigTransition, mixed) => mixed>,

50 };

52 // TODO: Is there a way to not include the tag or name here?

packages/react-reconciler/src/ReactFiberTransition.js

+43 -4

@@ -7,12 +7,20 @@

7 * @flow

8 */

9 import type {Fiber, FiberRoot} from './ReactInternalTypes';

10 +import type {Thenable} from 'shared/ReactTypes';

11 import type {Lanes} from './ReactFiberLane';

12 import type {StackCursor} from './ReactFiberStack';

13 import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent';

13 -import type {Transition} from './ReactFiberTracingMarkerComponent';

14 +import type {

15 + BatchConfigTransition,

16 + Transition,

17 +} from './ReactFiberTracingMarkerComponent';

15 -import {enableCache, enableTransitionTracing} from 'shared/ReactFeatureFlags';

19 +import {

20 + enableCache,

21 + enableTransitionTracing,

22 + enableAsyncActions,

23 +} from 'shared/ReactFeatureFlags';

24 import {isPrimaryRenderer} from './ReactFiberConfig';

25 import {createCursor, push, pop} from './ReactFiberStack';

26 import {

@@ -26,13 +34,44 @@ import {

34 } from './ReactFiberCacheComponent';

36 import ReactSharedInternals from 'shared/ReactSharedInternals';

37 +import {entangleAsyncAction} from './ReactFiberAsyncAction';

39 const {ReactCurrentBatchConfig} = ReactSharedInternals;

41 export const NoTransition = null;

34 -export function requestCurrentTransition(): Transition | null {

35 - return ReactCurrentBatchConfig.transition;

43 +export function requestCurrentTransition(): BatchConfigTransition | null {

44 + const transition = ReactCurrentBatchConfig.transition;

45 + if (transition !== null) {

46 + // Whenever a transition update is scheduled, register a callback on the

47 + // transition object so we can get the return value of the scope function.

48 + transition._callbacks.add(handleTransitionScopeResult);

49 + }

50 + return transition;

51 +}

52 +

53 +function handleTransitionScopeResult(

54 + transition: BatchConfigTransition,

55 + returnValue: mixed,

56 +): void {

57 + if (

58 + enableAsyncActions &&

59 + returnValue !== null &&

60 + typeof returnValue === 'object' &&

61 + typeof returnValue.then === 'function'

62 + ) {

63 + // This is an async action.

64 + const thenable: Thenable<mixed> = (returnValue: any);

65 + entangleAsyncAction(transition, thenable);

66 + }

67 +}

68 +

69 +export function notifyTransitionCallbacks(

70 + transition: BatchConfigTransition,

71 + returnValue: mixed,

72 +) {

73 + const callbacks = transition._callbacks;

74 + callbacks.forEach(callback => callback(transition, returnValue));

75 }

77 // When retrying a Suspense/Offscreen boundary, we restore the cache that was

packages/react-reconciler/src/ReactFiberWorkLoop.js

+11 -10

@@ -161,6 +161,7 @@ import {

161 OffscreenLane,

162 SyncUpdateLanes,

163 UpdateLanes,

164 + claimNextTransitionLane,

165 } from './ReactFiberLane';

166 import {

167 DiscreteEventPriority,

@@ -170,7 +171,7 @@ import {

171 lowerEventPriority,

172 lanesToEventPriority,

173 } from './ReactEventPriorities';

173 -import {requestCurrentTransition, NoTransition} from './ReactFiberTransition';

174 +import {requestCurrentTransition} from './ReactFiberTransition';

175 import {

176 SelectiveHydrationException,

177 beginWork as originalBeginWork,

@@ -633,15 +634,15 @@ export function requestUpdateLane(fiber: Fiber): Lane {

634 return pickArbitraryLane(workInProgressRootRenderLanes);

635 }

636

636 - const isTransition = requestCurrentTransition() !== NoTransition;

637 - if (isTransition) {

638 - if (__DEV__ && ReactCurrentBatchConfig.transition !== null) {

639 - const transition = ReactCurrentBatchConfig.transition;

640 - if (!transition._updatedFibers) {

641 - transition._updatedFibers = new Set();

637 + const transition = requestCurrentTransition();

638 + if (transition !== null) {

639 + if (__DEV__) {

640 + const batchConfigTransition = ReactCurrentBatchConfig.transition;

641 + if (!batchConfigTransition._updatedFibers) {

642 + batchConfigTransition._updatedFibers = new Set();

643 }

644

644 - transition._updatedFibers.add(fiber);

645 + batchConfigTransition._updatedFibers.add(fiber);

646 }

647

648 const actionScopeLane = peekEntangledActionLane();

@@ -651,7 +652,7 @@ export function requestUpdateLane(fiber: Fiber): Lane {

652 : // We may or may not be inside an async action scope. If we are, this

653 // is the first update in that scope. Either way, we need to get a

654 // fresh transition lane.

654 - requestTransitionLane();

655 + requestTransitionLane(transition);

656 }

657

658 // Updates originating inside certain React methods, like flushSync, have

@@ -712,7 +713,7 @@ export function requestDeferredLane(): Lane {

713 workInProgressDeferredLane = OffscreenLane;

714 } else {

715 // Everything else is spawned as a transition.

715 - workInProgressDeferredLane = requestTransitionLane();

716 + workInProgressDeferredLane = claimNextTransitionLane();

717 }

718 }

719

packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js

+85

@@ -1641,4 +1641,89 @@ describe('ReactAsyncActions', () => {

1641 expect(root).toMatchRenderedOutput(D);

1642 },

1643 );

1644 +

1645 + // @gate enableAsyncActions

1646 + test('React.startTransition supports async actions', async () => {

1647 + const startTransition = React.startTransition;

1648 +

1649 + function App({text}) {

1650 + return <Text text={text} />;

1651 + }

1652 +

1653 + const root = ReactNoop.createRoot();

1654 + await act(() => {

1655 + root.render(<App text="A" />);

1656 + });

1657 + assertLog(['A']);

1658 +

1659 + await act(() => {

1660 + startTransition(async () => {

1661 + // Update to B

1662 + root.render(<App text="B" />);

1663 +

1664 + // There's an async gap before C is updated

1665 + await getText('Wait before updating to C');

1666 + root.render(<App text="C" />);

1667 +

1668 + Scheduler.log('Async action ended');

1669 + });

1670 + });

1671 + // The update to B is blocked because the async action hasn't completed yet.

1672 + assertLog([]);

1673 + expect(root).toMatchRenderedOutput('A');

1674 +

1675 + // Finish the async action

1676 + await act(() => resolveText('Wait before updating to C'));

1677 +

1678 + // Now both B and C can finish in a single batch.

1679 + assertLog(['Async action ended', 'C']);

1680 + expect(root).toMatchRenderedOutput('C');

1681 + });

1682 +

1683 + // @gate enableAsyncActions

1684 + test('useOptimistic works with async actions passed to React.startTransition', async () => {

1685 + const startTransition = React.startTransition;

1686 +

1687 + let setOptimisticText;

1688 + function App({text: canonicalText}) {

1689 + const [text, _setOptimisticText] = useOptimistic(

1690 + canonicalText,

1691 + (_, optimisticText) => `${optimisticText} (loading...)`,

1692 + );

1693 + setOptimisticText = _setOptimisticText;

1694 + return (

1695 +

1696 + <Text text={text} />

1697 +

1698 + );

1699 + }

1700 +

1701 + const root = ReactNoop.createRoot();

1702 + await act(() => {

1703 + root.render(<App text="Initial" />);

1704 + });

1705 + assertLog(['Initial']);

1706 + expect(root).toMatchRenderedOutput(Initial);

1707 +

1708 + // Start an async action using the non-hook form of startTransition. The

1709 + // action includes an optimistic update.

1710 + await act(() => {

1711 + startTransition(async () => {

1712 + Scheduler.log('Async action started');

1713 + setOptimisticText('Updated');

1714 + await getText('Yield before updating');

1715 + Scheduler.log('Async action ended');

1716 + startTransition(() => root.render(<App text="Updated" />));

1717 + });

1718 + });

1719 + // Because the action hasn't finished yet, the optimistic UI is shown.

1720 + assertLog(['Async action started', 'Updated (loading...)']);

1721 + expect(root).toMatchRenderedOutput(Updated (loading...));

1722 +

1723 + // Finish the async action. The optimistic state is reverted and replaced by

1724 + // the canonical state.

1725 + await act(() => resolveText('Yield before updating'));

1726 + assertLog(['Async action ended', 'Updated']);

1727 + expect(root).toMatchRenderedOutput(Updated);

1728 + });

1729 });

packages/react/src/ReactStartTransition.js

+9 -2

@@ -17,7 +17,13 @@ export function startTransition(

17 options?: StartTransitionOptions,

18 ) {

19 const prevTransition = ReactCurrentBatchConfig.transition;

20 - ReactCurrentBatchConfig.transition = ({}: BatchConfigTransition);

20 + // Each renderer registers a callback to receive the return value of

21 + // the scope function. This is used to implement async actions.

22 + const callbacks = new Set<(BatchConfigTransition, mixed) => mixed>();

23 + const transition: BatchConfigTransition = {

24 + _callbacks: callbacks,

25 + };

26 + ReactCurrentBatchConfig.transition = transition;

27 const currentTransition = ReactCurrentBatchConfig.transition;

29 if (__DEV__) {

@@ -34,7 +40,8 @@ export function startTransition(

40 }

42 try {

37 - scope();

43 + const returnValue = scope();

44 + callbacks.forEach(callback => callback(currentTransition, returnValue));

45 } finally {

46 ReactCurrentBatchConfig.transition = prevTransition;