@samitouri / QOS-React-2 / commits / ee5c194930

Fix async batching in React.startTransition (#29226)

Note: Despite the similar-sounding description, this fix is unrelated to the issue where updates that occur after an `await` in an async action must also be wrapped in their own `startTransition`, due to the absence of an AsyncContext mechanism in browsers today. --- Discovered a flaw in the current implementation of the isomorphic startTransition implementation (React.startTransition), related to async actions. It only creates an async scope if something calls setState within the synchronous part of the action (i.e. before the first `await`). I had thought this was fine because if there's no update during this part, then there's nothing that needs to be entangled. I didn't think this through, though — if there are multiple async updates interleaved throughout the rest of the action, we need the async scope to have already been created so that _those_ are batched together. An even easier way to observe this is to schedule an optimistic update after an `await` — the optimistic update should not be reverted until the async action is complete. To implement, during the reconciler's module initialization, we compose its startTransition implementation with any previous reconciler's startTransition that was already initialized. Then, the isomorphic startTransition is the composition of every reconciler's startTransition. ```js function startTransition(fn) { return startTransitionDOM(() => { return startTransitionART(() => { return startTransitionThreeFiber(() => { // and so on... return fn(); }); }); }); } ``` This is basically how flushSync is implemented, too.

Andrew Clark committed May 23, 2024 at 17:19 UTC ee5c19493086fdeb32057e16d1e3414370242307

7 files changed +269 -42

packages/react-dom/src/__tests__/ReactStartTransitionMultipleRenderers-test.js new

+143

@@ -0,0 +1,143 @@

1 +/**

2 + * Copyright (c) Meta Platforms, Inc. and affiliates.

3 + *

4 + * This source code is licensed under the MIT license found in the

5 + * LICENSE file in the root directory of this source tree.

6 + *

7 + * @emails react-core

8 + */

9 +

10 +'use strict';

11 +

12 +describe('ReactStartTransitionMultipleRenderers', () => {

13 + let act;

14 + let container;

15 + let React;

16 + let ReactDOMClient;

17 + let Scheduler;

18 + let assertLog;

19 + let startTransition;

20 + let useOptimistic;

21 + let textCache;

22 +

23 + beforeEach(() => {

24 + jest.resetModules();

25 + React = require('react');

26 + ReactDOMClient = require('react-dom/client');

27 + Scheduler = require('scheduler');

28 + act = require('internal-test-utils').act;

29 + assertLog = require('internal-test-utils').assertLog;

30 + startTransition = React.startTransition;

31 + useOptimistic = React.useOptimistic;

32 + container = document.createElement('div');

33 + document.body.appendChild(container);

34 +

35 + textCache = new Map();

36 + });

37 +

38 + function resolveText(text) {

39 + const record = textCache.get(text);

40 + if (record === undefined) {

41 + const newRecord = {

42 + status: 'resolved',

43 + value: text,

44 + };

45 + textCache.set(text, newRecord);

46 + } else if (record.status === 'pending') {

47 + const thenable = record.value;

48 + record.status = 'resolved';

49 + record.value = text;

50 + thenable.pings.forEach(t => t(text));

51 + }

52 + }

53 +

54 + function getText(text) {

55 + const record = textCache.get(text);

56 + if (record === undefined) {

57 + const thenable = {

58 + pings: [],

59 + then(resolve) {

60 + if (newRecord.status === 'pending') {

61 + thenable.pings.push(resolve);

62 + } else {

63 + Promise.resolve().then(() => resolve(newRecord.value));

64 + }

65 + },

66 + };

67 + const newRecord = {

68 + status: 'pending',

69 + value: thenable,

70 + };

71 + textCache.set(text, newRecord);

72 + return thenable;

73 + } else {

74 + switch (record.status) {

75 + case 'pending':

76 + return record.value;

77 + case 'rejected':

78 + return Promise.reject(record.value);

79 + case 'resolved':

80 + return Promise.resolve(record.value);

81 + }

82 + }

83 + }

84 +

85 + function Text({text}) {

86 + Scheduler.log(text);

87 + return text;

88 + }

89 +

90 + afterEach(() => {

91 + document.body.removeChild(container);

92 + });

93 +

94 + // This test imports multiple reconcilers. Because of how the renderers are

95 + // built, it only works when running tests using the actual build artifacts,

96 + // not the source files.

97 + // @gate !source

98 + test('React.startTransition works across multiple renderers', async () => {

99 + const ReactNoop = require('react-noop-renderer');

100 +

101 + const setIsPendings = new Set();

102 +

103 + function App() {

104 + const [isPending, setIsPending] = useOptimistic(false);

105 + setIsPendings.add(setIsPending);

106 + return <Text text={isPending ? 'Pending' : 'Not pending'} />;

107 + }

108 +

109 + const noopRoot = ReactNoop.createRoot(null);

110 + const domRoot = ReactDOMClient.createRoot(container);

111 +

112 + // Run the same component in two separate renderers.

113 + await act(() => {

114 + noopRoot.render(<App />);

115 + domRoot.render(<App />);

116 + });

117 + assertLog(['Not pending', 'Not pending']);

118 + expect(container.textContent).toEqual('Not pending');

119 + expect(noopRoot).toMatchRenderedOutput('Not pending');

120 +

121 + await act(() => {

122 + startTransition(async () => {

123 + // Wait until after an async gap before setting the optimistic state.

124 + await getText('Wait before setting isPending');

125 + setIsPendings.forEach(setIsPending => setIsPending(true));

126 +

127 + // The optimistic state should not be reverted until the

128 + // action completes.

129 + await getText('Wait until end of async action');

130 + });

131 + });

132 +

133 + await act(() => resolveText('Wait before setting isPending'));

134 + assertLog(['Pending', 'Pending']);

135 + expect(container.textContent).toEqual('Pending');

136 + expect(noopRoot).toMatchRenderedOutput('Pending');

137 +

138 + await act(() => resolveText('Wait until end of async action'));

139 + assertLog(['Not pending', 'Not pending']);

140 + expect(container.textContent).toEqual('Not pending');

141 + expect(noopRoot).toMatchRenderedOutput('Not pending');

142 + });

143 +});

packages/react-reconciler/src/ReactFiberHooks.js

+11 -12

@@ -154,10 +154,7 @@ import {

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

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

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

157 -import {

158 - notifyTransitionCallbacks,

159 - requestCurrentTransition,

160 -} from './ReactFiberTransition';

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

158

159 export type Update<S, A> = {

160 lane: Lane,

@@ -2020,9 +2017,7 @@ function runActionStateAction<S, P>(

2017

2018 // This is a fork of startTransition

2019 const prevTransition = ReactSharedInternals.T;

2023 - const currentTransition: BatchConfigTransition = {

2024 - _callbacks: new Set<(BatchConfigTransition, mixed) => mixed>(),

2025 - };

2020 + const currentTransition: BatchConfigTransition = {};

2021 ReactSharedInternals.T = currentTransition;

2022 if (__DEV__) {

2023 ReactSharedInternals.T._updatedFibers = new Set();

@@ -2034,6 +2029,10 @@ function runActionStateAction<S, P>(

2029

2030 try {

2031 const returnValue = action(prevState, payload);

2032 + const onStartTransitionFinish = ReactSharedInternals.S;

2033 + if (onStartTransitionFinish !== null) {

2034 + onStartTransitionFinish(currentTransition, returnValue);

2035 + }

2036 if (

2037 returnValue !== null &&

2038 typeof returnValue === 'object' &&

@@ -2041,7 +2040,6 @@ function runActionStateAction<S, P>(

2040 typeof returnValue.then === 'function'

2041 ) {

2042 const thenable = ((returnValue: any): Thenable<Awaited<S>>);

2044 - notifyTransitionCallbacks(currentTransition, thenable);

2043

2044 // Attach a listener to read the return state of the action. As soon as

2045 // this resolves, we can run the next action in the sequence.

@@ -2843,9 +2841,7 @@ function startTransition<S>(

2841 );

2842

2843 const prevTransition = ReactSharedInternals.T;

2846 - const currentTransition: BatchConfigTransition = {

2847 - _callbacks: new Set<(BatchConfigTransition, mixed) => mixed>(),

2848 - };

2844 + const currentTransition: BatchConfigTransition = {};

2845

2846 if (enableAsyncActions) {

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

@@ -2876,6 +2872,10 @@ function startTransition<S>(

2872 try {

2873 if (enableAsyncActions) {

2874 const returnValue = callback();

2875 + const onStartTransitionFinish = ReactSharedInternals.S;

2876 + if (onStartTransitionFinish !== null) {

2877 + onStartTransitionFinish(currentTransition, returnValue);

2878 + }

2879

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

2881 // this new action with the existing scope.

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

2891 typeof returnValue.then === 'function'

2892 ) {

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

2894 - notifyTransitionCallbacks(currentTransition, thenable);

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

2895 // action has completed.

2896 const thenableForFinishedState = chainThenableValue(

packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js

-1

@@ -47,7 +47,6 @@ export type BatchConfigTransition = {

47 name?: string,

48 startTime?: number,

49 _updatedFibers?: Set<Fiber>,

50 - _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

+38 -22

@@ -38,32 +38,48 @@ import {entangleAsyncAction} from './ReactFiberAsyncAction';

39 export const NoTransition = null;

41 -export function requestCurrentTransition(): BatchConfigTransition | null {

42 - const transition = ReactSharedInternals.T;

43 - if (transition !== null) {

44 - // Whenever a transition update is scheduled, register a callback on the

45 - // transition object so we can get the return value of the scope function.

46 - transition._callbacks.add((handleAsyncAction: any));

47 - }

48 - return transition;

49 -}

50 -

51 -function handleAsyncAction(

41 +// Attach this reconciler instance's onStartTransitionFinish implementation to

42 +// the shared internals object. This is used by the isomorphic implementation of

43 +// startTransition to compose all the startTransitions together.

44 +//

45 +// function startTransition(fn) {

46 +// return startTransitionDOM(() => {

47 +// return startTransitionART(() => {

48 +// return startTransitionThreeFiber(() => {

49 +// // and so on...

50 +// return fn();

51 +// });

52 +// });

53 +// });

54 +// }

55 +//

56 +// Currently we only compose together the code that runs at the end of each

57 +// startTransition, because for now that's sufficient — the part that sets

58 +// isTransition=true on the stack uses a separate shared internal field. But

59 +// really we should delete the shared field and track isTransition per

60 +// reconciler. Leaving this for a future PR.

61 +const prevOnStartTransitionFinish = ReactSharedInternals.S;

62 +ReactSharedInternals.S = function onStartTransitionFinishForReconciler(

63 transition: BatchConfigTransition,

53 - thenable: Thenable<mixed>,

54 -): void {

55 - if (enableAsyncActions) {

56 - // This is an async action.

64 + returnValue: mixed,

65 +) {

66 + if (

67 + enableAsyncActions &&

68 + typeof returnValue === 'object' &&

69 + returnValue !== null &&

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

71 + ) {

72 + // This is an async action

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

74 entangleAsyncAction(transition, thenable);

75 }

59 -}

76 + if (prevOnStartTransitionFinish !== null) {

77 + prevOnStartTransitionFinish(transition, returnValue);

78 + }

79 +};

61 -export function notifyTransitionCallbacks(

62 - transition: BatchConfigTransition,

63 - returnValue: mixed,

64 -) {

65 - const callbacks = transition._callbacks;

66 - callbacks.forEach(callback => callback(transition, returnValue));

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

82 + return ReactSharedInternals.T;

83 }

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

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

+70

@@ -1731,6 +1731,76 @@ describe('ReactAsyncActions', () => {

1731 expect(root).toMatchRenderedOutput(Updated);

1732 });

1733

1734 + // @gate enableAsyncActions

1735 + test(

1736 + 'regression: updates in an action passed to React.startTransition are batched ' +

1737 + 'even if there were no updates before the first await',

1738 + async () => {

1739 + // Regression for a bug that occured in an older, too-clever-by-half

1740 + // implementation of the isomorphic startTransition API. Now, the

1741 + // isomorphic startTransition is literally the composition of every

1742 + // reconciler instance's startTransition, so the behavior is less likely

1743 + // to regress in the future.

1744 + const startTransition = React.startTransition;

1745 +

1746 + let setOptimisticText;

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

1748 + const [text, _setOptimisticText] = useOptimistic(

1749 + canonicalText,

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

1751 + );

1752 + setOptimisticText = _setOptimisticText;

1753 + return (

1754 +

1755 + <Text text={text} />

1756 +

1757 + );

1758 + }

1759 +

1760 + const root = ReactNoop.createRoot();

1761 + await act(() => {

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

1763 + });

1764 + assertLog(['Initial']);

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

1766 +

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

1768 + // action includes an optimistic update.

1769 + await act(() => {

1770 + startTransition(async () => {

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

1772 +

1773 + // Yield to an async task *before* any updates have occurred.

1774 + await getText('Yield before optimistic update');

1775 +

1776 + // This optimistic update happens after an async gap. In the

1777 + // regression case, this update was not correctly associated with

1778 + // the outer async action, causing the optimistic update to be

1779 + // immediately reverted.

1780 + setOptimisticText('Updated');

1781 +

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

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

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

1785 + });

1786 + });

1787 + assertLog(['Async action started']);

1788 +

1789 + // Wait for an async gap, then schedule an optimistic update.

1790 + await act(() => resolveText('Yield before optimistic update'));

1791 +

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

1793 + assertLog(['Updated (loading...)']);

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

1795 +

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

1797 + // by the canonical state.

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

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

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

1801 + },

1802 + );

1803 +

1804 test('React.startTransition captures async errors and passes them to reportError', async () => {

1805 // NOTE: This is gated here instead of using the pragma because the failure

1806 // happens asynchronously and the `gate` runtime doesn't capture it.

packages/react/src/ReactSharedInternalsClient.js

@@ -15,6 +15,7 @@ export type SharedStateClient = {

15 H: null | Dispatcher, // ReactCurrentDispatcher for Hooks

16 A: null | AsyncDispatcher, // ReactCurrentCache for Cache

17 T: null | BatchConfigTransition, // ReactCurrentBatchConfig for Transitions

18 + S: null | ((BatchConfigTransition, mixed) => void), // onStartTransitionFinish

20 // DEV-only

@@ -43,6 +44,7 @@ const ReactSharedInternals: SharedStateClient = ({

44 H: null,

45 A: null,

46 T: null,

47 + S: null,

48 }: any);

50 if (__DEV__) {

packages/react/src/ReactStartTransition.js

+5 -7

@@ -23,12 +23,7 @@ export function startTransition(

23 options?: StartTransitionOptions,

24 ) {

25 const prevTransition = ReactSharedInternals.T;

26 - // Each renderer registers a callback to receive the return value of

27 - // the scope function. This is used to implement async actions.

28 - const callbacks = new Set<(BatchConfigTransition, mixed) => mixed>();

29 - const transition: BatchConfigTransition = {

30 - _callbacks: callbacks,

31 - };

26 + const transition: BatchConfigTransition = {};

27 ReactSharedInternals.T = transition;

28 const currentTransition = ReactSharedInternals.T;

@@ -48,12 +43,15 @@ export function startTransition(

43 if (enableAsyncActions) {

44 try {

45 const returnValue = scope();

46 + const onStartTransitionFinish = ReactSharedInternals.S;

47 + if (onStartTransitionFinish !== null) {

48 + onStartTransitionFinish(transition, returnValue);

49 + }

50 if (

51 typeof returnValue === 'object' &&

52 returnValue !== null &&

53 typeof returnValue.then === 'function'

54 ) {

56 - callbacks.forEach(callback => callback(currentTransition, returnValue));

55 returnValue.then(noop, reportGlobalError);

56 }

57 } catch (error) {