From f032dfe2845ff2895edaec40c60339f9600eca54 Mon Sep 17 00:00:00 2001 From: Andrew Clark Date: Wed, 24 Jan 2024 18:02:10 -0500 Subject: [PATCH] Batch async updates to classes, roots too This implements the previous fix for class components and root-level updates, too. I put this in its own step because there are some extra changes related to avoiding a context stack mismatch. --- .../src/ReactFiberBeginWork.js | 7 + .../src/ReactFiberClassComponent.js | 4 + .../src/ReactFiberClassUpdateQueue.js | 37 ++++ .../src/__tests__/ReactAsyncActions-test.js | 180 ++++++++++++++++++ 4 files changed, 228 insertions(+) diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js index f2a96b4d34d9f..5c555b1812d43 100644 --- a/packages/react-reconciler/src/ReactFiberBeginWork.js +++ b/packages/react-reconciler/src/ReactFiberBeginWork.js @@ -137,6 +137,7 @@ import { cloneUpdateQueue, initializeUpdateQueue, enqueueCapturedUpdate, + suspendIfUpdateReadFromEntangledAsyncAction, } from './ReactFiberClassUpdateQueue'; import { NoLane, @@ -945,6 +946,7 @@ function updateCacheComponent( if (includesSomeLane(current.lanes, renderLanes)) { cloneUpdateQueue(current, workInProgress); processUpdateQueue(workInProgress, null, null, renderLanes); + suspendIfUpdateReadFromEntangledAsyncAction(); } const prevState: CacheComponentState = current.memoizedState; const nextState: CacheComponentState = workInProgress.memoizedState; @@ -1475,6 +1477,11 @@ function updateHostRoot( } } + // This would ideally go inside processUpdateQueue, but because it suspends, + // it needs to happen after the `pushCacheProvider` call above to avoid a + // context stack mismatch. A bit unfortunate. + suspendIfUpdateReadFromEntangledAsyncAction(); + // Caution: React DevTools currently depends on this property // being called "element". const nextChildren = nextState.element; diff --git a/packages/react-reconciler/src/ReactFiberClassComponent.js b/packages/react-reconciler/src/ReactFiberClassComponent.js index c816fb0ba21ea..cceb78c8ed878 100644 --- a/packages/react-reconciler/src/ReactFiberClassComponent.js +++ b/packages/react-reconciler/src/ReactFiberClassComponent.js @@ -53,6 +53,7 @@ import { ForceUpdate, initializeUpdateQueue, cloneUpdateQueue, + suspendIfUpdateReadFromEntangledAsyncAction, } from './ReactFiberClassUpdateQueue'; import {NoLanes} from './ReactFiberLane'; import { @@ -892,6 +893,7 @@ function mountClassInstance( // If we had additional state updates during this life-cycle, let's // process them now. processUpdateQueue(workInProgress, newProps, instance, renderLanes); + suspendIfUpdateReadFromEntangledAsyncAction(); instance.state = workInProgress.memoizedState; } @@ -959,6 +961,7 @@ function resumeMountClassInstance( const oldState = workInProgress.memoizedState; let newState = (instance.state = oldState); processUpdateQueue(workInProgress, newProps, instance, renderLanes); + suspendIfUpdateReadFromEntangledAsyncAction(); newState = workInProgress.memoizedState; if ( oldProps === newProps && @@ -1109,6 +1112,7 @@ function updateClassInstance( const oldState = workInProgress.memoizedState; let newState = (instance.state = oldState); processUpdateQueue(workInProgress, newProps, instance, renderLanes); + suspendIfUpdateReadFromEntangledAsyncAction(); newState = workInProgress.memoizedState; if ( diff --git a/packages/react-reconciler/src/ReactFiberClassUpdateQueue.js b/packages/react-reconciler/src/ReactFiberClassUpdateQueue.js index 2e88c982022d3..d7a1efbb47f4a 100644 --- a/packages/react-reconciler/src/ReactFiberClassUpdateQueue.js +++ b/packages/react-reconciler/src/ReactFiberClassUpdateQueue.js @@ -125,6 +125,10 @@ import { import {setIsStrictModeForDevtools} from './ReactFiberDevToolsHook'; import assign from 'shared/assign'; +import { + peekEntangledActionLane, + peekEntangledActionThenable, +} from './ReactFiberAsyncAction'; export type Update = { lane: Lane, @@ -463,12 +467,38 @@ function getStateFromUpdate( return prevState; } +let didReadFromEntangledAsyncAction: boolean = false; + +// Each call to processUpdateQueue should be accompanied by a call to this. It's +// only in a separate function because in updateHostRoot, it must happen after +// all the context stacks have been pushed to, to prevent a stack mismatch. A +// bit unfortunate. +export function suspendIfUpdateReadFromEntangledAsyncAction() { + // 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. + // TODO: Once we support hooks inside useMemo (or an equivalent + // memoization boundary like Forget), hoist this logic so that it only + // suspends if the memo boundary produces a new value. + if (didReadFromEntangledAsyncAction) { + const entangledActionThenable = peekEntangledActionThenable(); + if (entangledActionThenable !== null) { + // TODO: Instead of the throwing the thenable directly, throw a + // special object like `use` does so we can detect if it's captured + // by userspace. + throw entangledActionThenable; + } + } +} + export function processUpdateQueue( workInProgress: Fiber, props: any, instance: any, renderLanes: Lanes, ): void { + didReadFromEntangledAsyncAction = false; + // This is always non-null on a ClassComponent or HostRoot const queue: UpdateQueue = (workInProgress.updateQueue: any); @@ -571,6 +601,13 @@ export function processUpdateQueue( } 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; + } + if (newLastBaseUpdate !== null) { const clone: Update = { // This update is going to be committed so we never want uncommit diff --git a/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js b/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js index 4d4d8558ff5aa..1f45fd84d0430 100644 --- a/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js +++ b/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js @@ -1461,4 +1461,184 @@ describe('ReactAsyncActions', () => { expect(root).toMatchRenderedOutput(C); }, ); + + // @gate enableAsyncActions + test( + 'updates in an async action are entangled even if useTransition hook ' + + 'is unmounted before it finishes (class component)', + async () => { + let startTransition; + function Updater() { + const [isPending, _start] = useTransition(); + startTransition = _start; + return ( + + + + ); + } + + let setText; + class Sibling extends React.Component { + state = {text: 'A'}; + render() { + setText = text => this.setState({text}); + return ( + + + + ); + } + } + + function App({showUpdater}) { + return ( + <> + {showUpdater ? : null} + + + ); + } + + const root = ReactNoop.createRoot(); + await act(() => { + root.render(); + }); + assertLog(['Pending: false', 'A']); + expect(root).toMatchRenderedOutput( + <> + Pending: false + A + , + ); + + // Start an async action that has multiple updates with async + // operations in between. + await act(() => { + startTransition(async () => { + Scheduler.log('Async action started'); + startTransition(() => setText('B')); + + await getText('Wait before updating to C'); + + Scheduler.log('Async action ended'); + startTransition(() => setText('C')); + }); + }); + assertLog(['Async action started', 'Pending: true']); + expect(root).toMatchRenderedOutput( + <> + Pending: true + A + , + ); + + // Delete the component that contains the useTransition hook. This + // component no longer blocks the transition from completing. But the + // pending update to Sibling should not be allowed to finish, because it's + // part of the async action. + await act(() => { + root.render(); + }); + assertLog(['A']); + expect(root).toMatchRenderedOutput(A); + + // Finish the async action. Notice the intermediate B state was never + // shown, because it was batched with the update that came later in the + // same action. + await act(() => resolveText('Wait before updating to C')); + assertLog(['Async action ended', 'C']); + expect(root).toMatchRenderedOutput(C); + + // Check that subsequent updates are unaffected. + await act(() => setText('D')); + assertLog(['D']); + expect(root).toMatchRenderedOutput(D); + }, + ); + + // @gate enableAsyncActions + test( + 'updates in an async action are entangled even if useTransition hook ' + + 'is unmounted before it finishes (root update)', + async () => { + let startTransition; + function Updater() { + const [isPending, _start] = useTransition(); + startTransition = _start; + return ( + + + + ); + } + + let setShowUpdater; + function App({text}) { + const [showUpdater, _setShowUpdater] = useState(true); + setShowUpdater = _setShowUpdater; + return ( + <> + {showUpdater ? : null} + + + + + ); + } + + const root = ReactNoop.createRoot(); + await act(() => { + root.render(); + }); + assertLog(['Pending: false', 'A']); + expect(root).toMatchRenderedOutput( + <> + Pending: false + A + , + ); + + // Start an async action that has multiple updates with async + // operations in between. + await act(() => { + startTransition(async () => { + Scheduler.log('Async action started'); + startTransition(() => root.render()); + + await getText('Wait before updating to C'); + + Scheduler.log('Async action ended'); + startTransition(() => root.render()); + }); + }); + assertLog(['Async action started', 'Pending: true']); + expect(root).toMatchRenderedOutput( + <> + Pending: true + A + , + ); + + // Delete the component that contains the useTransition hook. This + // component no longer blocks the transition from completing. But the + // pending update to Sibling should not be allowed to finish, because it's + // part of the async action. + await act(() => setShowUpdater(false)); + assertLog(['A']); + expect(root).toMatchRenderedOutput(A); + + // Finish the async action. Notice the intermediate B state was never + // shown, because it was batched with the update that came later in the + // same action. + await act(() => resolveText('Wait before updating to C')); + assertLog(['Async action ended', 'C']); + expect(root).toMatchRenderedOutput(C); + + // Check that subsequent updates are unaffected. + await act(() => root.render()); + assertLog(['D']); + expect(root).toMatchRenderedOutput(D); + }, + ); });