From ee5c19493086fdeb32057e16d1e3414370242307 Mon Sep 17 00:00:00 2001 From: Andrew Clark Date: Thu, 23 May 2024 17:19:09 -0400 Subject: [PATCH] Fix async batching in React.startTransition (#29226) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- ...ctStartTransitionMultipleRenderers-test.js | 143 ++++++++++++++++++ .../react-reconciler/src/ReactFiberHooks.js | 23 ++- .../src/ReactFiberTracingMarkerComponent.js | 1 - .../src/ReactFiberTransition.js | 60 +++++--- .../src/__tests__/ReactAsyncActions-test.js | 70 +++++++++ .../react/src/ReactSharedInternalsClient.js | 2 + packages/react/src/ReactStartTransition.js | 12 +- 7 files changed, 269 insertions(+), 42 deletions(-) create mode 100644 packages/react-dom/src/__tests__/ReactStartTransitionMultipleRenderers-test.js diff --git a/packages/react-dom/src/__tests__/ReactStartTransitionMultipleRenderers-test.js b/packages/react-dom/src/__tests__/ReactStartTransitionMultipleRenderers-test.js new file mode 100644 index 000000000000..5318cc91fb26 --- /dev/null +++ b/packages/react-dom/src/__tests__/ReactStartTransitionMultipleRenderers-test.js @@ -0,0 +1,143 @@ +/** + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + * + * @emails react-core + */ + +'use strict'; + +describe('ReactStartTransitionMultipleRenderers', () => { + let act; + let container; + let React; + let ReactDOMClient; + let Scheduler; + let assertLog; + let startTransition; + let useOptimistic; + let textCache; + + beforeEach(() => { + jest.resetModules(); + React = require('react'); + ReactDOMClient = require('react-dom/client'); + Scheduler = require('scheduler'); + act = require('internal-test-utils').act; + assertLog = require('internal-test-utils').assertLog; + startTransition = React.startTransition; + useOptimistic = React.useOptimistic; + container = document.createElement('div'); + document.body.appendChild(container); + + textCache = new Map(); + }); + + function resolveText(text) { + const record = textCache.get(text); + if (record === undefined) { + const newRecord = { + status: 'resolved', + value: text, + }; + textCache.set(text, newRecord); + } else if (record.status === 'pending') { + const thenable = record.value; + record.status = 'resolved'; + record.value = text; + thenable.pings.forEach(t => t(text)); + } + } + + function getText(text) { + const record = textCache.get(text); + if (record === undefined) { + const thenable = { + pings: [], + then(resolve) { + if (newRecord.status === 'pending') { + thenable.pings.push(resolve); + } else { + Promise.resolve().then(() => resolve(newRecord.value)); + } + }, + }; + const newRecord = { + status: 'pending', + value: thenable, + }; + textCache.set(text, newRecord); + return thenable; + } else { + switch (record.status) { + case 'pending': + return record.value; + case 'rejected': + return Promise.reject(record.value); + case 'resolved': + return Promise.resolve(record.value); + } + } + } + + function Text({text}) { + Scheduler.log(text); + return text; + } + + afterEach(() => { + document.body.removeChild(container); + }); + + // This test imports multiple reconcilers. Because of how the renderers are + // built, it only works when running tests using the actual build artifacts, + // not the source files. + // @gate !source + test('React.startTransition works across multiple renderers', async () => { + const ReactNoop = require('react-noop-renderer'); + + const setIsPendings = new Set(); + + function App() { + const [isPending, setIsPending] = useOptimistic(false); + setIsPendings.add(setIsPending); + return ; + } + + const noopRoot = ReactNoop.createRoot(null); + const domRoot = ReactDOMClient.createRoot(container); + + // Run the same component in two separate renderers. + await act(() => { + noopRoot.render(); + domRoot.render(); + }); + assertLog(['Not pending', 'Not pending']); + expect(container.textContent).toEqual('Not pending'); + expect(noopRoot).toMatchRenderedOutput('Not pending'); + + await act(() => { + startTransition(async () => { + // Wait until after an async gap before setting the optimistic state. + await getText('Wait before setting isPending'); + setIsPendings.forEach(setIsPending => setIsPending(true)); + + // The optimistic state should not be reverted until the + // action completes. + await getText('Wait until end of async action'); + }); + }); + + await act(() => resolveText('Wait before setting isPending')); + assertLog(['Pending', 'Pending']); + expect(container.textContent).toEqual('Pending'); + expect(noopRoot).toMatchRenderedOutput('Pending'); + + await act(() => resolveText('Wait until end of async action')); + assertLog(['Not pending', 'Not pending']); + expect(container.textContent).toEqual('Not pending'); + expect(noopRoot).toMatchRenderedOutput('Not pending'); + }); +}); diff --git a/packages/react-reconciler/src/ReactFiberHooks.js b/packages/react-reconciler/src/ReactFiberHooks.js index d49b8367d099..17ee55814a7d 100644 --- a/packages/react-reconciler/src/ReactFiberHooks.js +++ b/packages/react-reconciler/src/ReactFiberHooks.js @@ -154,10 +154,7 @@ import { import {HostTransitionContext} from './ReactFiberHostContext'; import {requestTransitionLane} from './ReactFiberRootScheduler'; import {isCurrentTreeHidden} from './ReactFiberHiddenContext'; -import { - notifyTransitionCallbacks, - requestCurrentTransition, -} from './ReactFiberTransition'; +import {requestCurrentTransition} from './ReactFiberTransition'; export type Update = { lane: Lane, @@ -2020,9 +2017,7 @@ function runActionStateAction( // This is a fork of startTransition const prevTransition = ReactSharedInternals.T; - const currentTransition: BatchConfigTransition = { - _callbacks: new Set<(BatchConfigTransition, mixed) => mixed>(), - }; + const currentTransition: BatchConfigTransition = {}; ReactSharedInternals.T = currentTransition; if (__DEV__) { ReactSharedInternals.T._updatedFibers = new Set(); @@ -2034,6 +2029,10 @@ function runActionStateAction( try { const returnValue = action(prevState, payload); + const onStartTransitionFinish = ReactSharedInternals.S; + if (onStartTransitionFinish !== null) { + onStartTransitionFinish(currentTransition, returnValue); + } if ( returnValue !== null && typeof returnValue === 'object' && @@ -2041,7 +2040,6 @@ function runActionStateAction( typeof returnValue.then === 'function' ) { const thenable = ((returnValue: any): Thenable>); - notifyTransitionCallbacks(currentTransition, thenable); // Attach a listener to read the return state of the action. As soon as // this resolves, we can run the next action in the sequence. @@ -2843,9 +2841,7 @@ function startTransition( ); const prevTransition = ReactSharedInternals.T; - const currentTransition: BatchConfigTransition = { - _callbacks: new Set<(BatchConfigTransition, mixed) => mixed>(), - }; + const currentTransition: BatchConfigTransition = {}; if (enableAsyncActions) { // We don't really need to use an optimistic update here, because we @@ -2876,6 +2872,10 @@ function startTransition( try { if (enableAsyncActions) { const returnValue = callback(); + const onStartTransitionFinish = ReactSharedInternals.S; + if (onStartTransitionFinish !== null) { + onStartTransitionFinish(currentTransition, returnValue); + } // Check if we're inside an async action scope. If so, we'll entangle // this new action with the existing scope. @@ -2891,7 +2891,6 @@ function startTransition( typeof returnValue.then === 'function' ) { const thenable = ((returnValue: any): Thenable); - notifyTransitionCallbacks(currentTransition, thenable); // Create a thenable that resolves to `finishedState` once the async // action has completed. const thenableForFinishedState = chainThenableValue( diff --git a/packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js b/packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js index 4a0da4958991..039daafa6fa1 100644 --- a/packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js +++ b/packages/react-reconciler/src/ReactFiberTracingMarkerComponent.js @@ -47,7 +47,6 @@ 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 df30ae5ddef3..c378beee30d3 100644 --- a/packages/react-reconciler/src/ReactFiberTransition.js +++ b/packages/react-reconciler/src/ReactFiberTransition.js @@ -38,32 +38,48 @@ import {entangleAsyncAction} from './ReactFiberAsyncAction'; export const NoTransition = null; -export function requestCurrentTransition(): BatchConfigTransition | null { - const transition = ReactSharedInternals.T; - 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((handleAsyncAction: any)); - } - return transition; -} - -function handleAsyncAction( +// Attach this reconciler instance's onStartTransitionFinish implementation to +// the shared internals object. This is used by the isomorphic implementation of +// startTransition to compose all the startTransitions together. +// +// function startTransition(fn) { +// return startTransitionDOM(() => { +// return startTransitionART(() => { +// return startTransitionThreeFiber(() => { +// // and so on... +// return fn(); +// }); +// }); +// }); +// } +// +// Currently we only compose together the code that runs at the end of each +// startTransition, because for now that's sufficient — the part that sets +// isTransition=true on the stack uses a separate shared internal field. But +// really we should delete the shared field and track isTransition per +// reconciler. Leaving this for a future PR. +const prevOnStartTransitionFinish = ReactSharedInternals.S; +ReactSharedInternals.S = function onStartTransitionFinishForReconciler( transition: BatchConfigTransition, - thenable: Thenable, -): void { - if (enableAsyncActions) { - // This is an async action. + returnValue: mixed, +) { + if ( + enableAsyncActions && + typeof returnValue === 'object' && + returnValue !== null && + typeof returnValue.then === 'function' + ) { + // This is an async action + const thenable: Thenable = (returnValue: any); entangleAsyncAction(transition, thenable); } -} + if (prevOnStartTransitionFinish !== null) { + prevOnStartTransitionFinish(transition, returnValue); + } +}; -export function notifyTransitionCallbacks( - transition: BatchConfigTransition, - returnValue: mixed, -) { - const callbacks = transition._callbacks; - callbacks.forEach(callback => callback(transition, returnValue)); +export function requestCurrentTransition(): BatchConfigTransition | null { + return ReactSharedInternals.T; } // When retrying a Suspense/Offscreen boundary, we restore the cache that was diff --git a/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js b/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js index 0a29be2b6fec..f7faf79aaa7d 100644 --- a/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js +++ b/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js @@ -1731,6 +1731,76 @@ describe('ReactAsyncActions', () => { expect(root).toMatchRenderedOutput(Updated); }); + // @gate enableAsyncActions + test( + 'regression: updates in an action passed to React.startTransition are batched ' + + 'even if there were no updates before the first await', + async () => { + // Regression for a bug that occured in an older, too-clever-by-half + // implementation of the isomorphic startTransition API. Now, the + // isomorphic startTransition is literally the composition of every + // reconciler instance's startTransition, so the behavior is less likely + // to regress in the future. + 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'); + + // Yield to an async task *before* any updates have occurred. + await getText('Yield before optimistic update'); + + // This optimistic update happens after an async gap. In the + // regression case, this update was not correctly associated with + // the outer async action, causing the optimistic update to be + // immediately reverted. + setOptimisticText('Updated'); + + await getText('Yield before updating'); + Scheduler.log('Async action ended'); + startTransition(() => root.render()); + }); + }); + assertLog(['Async action started']); + + // Wait for an async gap, then schedule an optimistic update. + await act(() => resolveText('Yield before optimistic update')); + + // Because the action hasn't finished yet, the optimistic UI is shown. + assertLog(['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); + }, + ); + test('React.startTransition captures async errors and passes them to reportError', async () => { // NOTE: This is gated here instead of using the pragma because the failure // happens asynchronously and the `gate` runtime doesn't capture it. diff --git a/packages/react/src/ReactSharedInternalsClient.js b/packages/react/src/ReactSharedInternalsClient.js index 3c4f47d94422..452bd933dab0 100644 --- a/packages/react/src/ReactSharedInternalsClient.js +++ b/packages/react/src/ReactSharedInternalsClient.js @@ -15,6 +15,7 @@ export type SharedStateClient = { H: null | Dispatcher, // ReactCurrentDispatcher for Hooks A: null | AsyncDispatcher, // ReactCurrentCache for Cache T: null | BatchConfigTransition, // ReactCurrentBatchConfig for Transitions + S: null | ((BatchConfigTransition, mixed) => void), // onStartTransitionFinish // DEV-only @@ -43,6 +44,7 @@ const ReactSharedInternals: SharedStateClient = ({ H: null, A: null, T: null, + S: null, }: any); if (__DEV__) { diff --git a/packages/react/src/ReactStartTransition.js b/packages/react/src/ReactStartTransition.js index 5311e1d29ad7..6bae8947e1ce 100644 --- a/packages/react/src/ReactStartTransition.js +++ b/packages/react/src/ReactStartTransition.js @@ -23,12 +23,7 @@ export function startTransition( options?: StartTransitionOptions, ) { const prevTransition = ReactSharedInternals.T; - // 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, - }; + const transition: BatchConfigTransition = {}; ReactSharedInternals.T = transition; const currentTransition = ReactSharedInternals.T; @@ -48,12 +43,15 @@ export function startTransition( if (enableAsyncActions) { try { const returnValue = scope(); + const onStartTransitionFinish = ReactSharedInternals.S; + if (onStartTransitionFinish !== null) { + onStartTransitionFinish(transition, returnValue); + } if ( typeof returnValue === 'object' && returnValue !== null && typeof returnValue.then === 'function' ) { - callbacks.forEach(callback => callback(currentTransition, returnValue)); returnValue.then(noop, reportGlobalError); } } catch (error) {