Skip to content

Commit f032dfe

Browse files
acdliteacdlite
committed
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.
1 parent 0ffbed7 commit f032dfe

4 files changed

Lines changed: 228 additions & 0 deletions

File tree

‎packages/react-reconciler/src/ReactFiberBeginWork.js‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ import {
137137
cloneUpdateQueue,
138138
initializeUpdateQueue,
139139
enqueueCapturedUpdate,
140+
suspendIfUpdateReadFromEntangledAsyncAction,
140141
} from './ReactFiberClassUpdateQueue';
141142
import {
142143
NoLane,
@@ -945,6 +946,7 @@ function updateCacheComponent(
945946
if (includesSomeLane(current.lanes, renderLanes)) {
946947
cloneUpdateQueue(current, workInProgress);
947948
processUpdateQueue(workInProgress, null, null, renderLanes);
949+
suspendIfUpdateReadFromEntangledAsyncAction();
948950
}
949951
const prevState: CacheComponentState = current.memoizedState;
950952
const nextState: CacheComponentState = workInProgress.memoizedState;
@@ -1475,6 +1477,11 @@ function updateHostRoot(
14751477
}
14761478
}
14771479

1480+
// This would ideally go inside processUpdateQueue, but because it suspends,
1481+
// it needs to happen after the `pushCacheProvider` call above to avoid a
1482+
// context stack mismatch. A bit unfortunate.
1483+
suspendIfUpdateReadFromEntangledAsyncAction();
1484+
14781485
// Caution: React DevTools currently depends on this property
14791486
// being called "element".
14801487
const nextChildren = nextState.element;

‎packages/react-reconciler/src/ReactFiberClassComponent.js‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ import {
5353
ForceUpdate,
5454
initializeUpdateQueue,
5555
cloneUpdateQueue,
56+
suspendIfUpdateReadFromEntangledAsyncAction,
5657
} from './ReactFiberClassUpdateQueue';
5758
import {NoLanes} from './ReactFiberLane';
5859
import {
@@ -892,6 +893,7 @@ function mountClassInstance(
892893
// If we had additional state updates during this life-cycle, let's
893894
// process them now.
894895
processUpdateQueue(workInProgress, newProps, instance, renderLanes);
896+
suspendIfUpdateReadFromEntangledAsyncAction();
895897
instance.state = workInProgress.memoizedState;
896898
}
897899

@@ -959,6 +961,7 @@ function resumeMountClassInstance(
959961
const oldState = workInProgress.memoizedState;
960962
let newState = (instance.state = oldState);
961963
processUpdateQueue(workInProgress, newProps, instance, renderLanes);
964+
suspendIfUpdateReadFromEntangledAsyncAction();
962965
newState = workInProgress.memoizedState;
963966
if (
964967
oldProps === newProps &&
@@ -1109,6 +1112,7 @@ function updateClassInstance(
11091112
const oldState = workInProgress.memoizedState;
11101113
let newState = (instance.state = oldState);
11111114
processUpdateQueue(workInProgress, newProps, instance, renderLanes);
1115+
suspendIfUpdateReadFromEntangledAsyncAction();
11121116
newState = workInProgress.memoizedState;
11131117

11141118
if (

‎packages/react-reconciler/src/ReactFiberClassUpdateQueue.js‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,10 @@ import {
125125
import {setIsStrictModeForDevtools} from './ReactFiberDevToolsHook';
126126

127127
import assign from 'shared/assign';
128+
import {
129+
peekEntangledActionLane,
130+
peekEntangledActionThenable,
131+
} from './ReactFiberAsyncAction';
128132

129133
export type Update<State> = {
130134
lane: Lane,
@@ -463,12 +467,38 @@ function getStateFromUpdate<State>(
463467
return prevState;
464468
}
465469

470+
let didReadFromEntangledAsyncAction: boolean = false;
471+
472+
// Each call to processUpdateQueue should be accompanied by a call to this. It's
473+
// only in a separate function because in updateHostRoot, it must happen after
474+
// all the context stacks have been pushed to, to prevent a stack mismatch. A
475+
// bit unfortunate.
476+
export function suspendIfUpdateReadFromEntangledAsyncAction() {
477+
// Check if this update is part of a pending async action. If so, we'll
478+
// need to suspend until the action has finished, so that it's batched
479+
// together with future updates in the same action.
480+
// TODO: Once we support hooks inside useMemo (or an equivalent
481+
// memoization boundary like Forget), hoist this logic so that it only
482+
// suspends if the memo boundary produces a new value.
483+
if (didReadFromEntangledAsyncAction) {
484+
const entangledActionThenable = peekEntangledActionThenable();
485+
if (entangledActionThenable !== null) {
486+
// TODO: Instead of the throwing the thenable directly, throw a
487+
// special object like `use` does so we can detect if it's captured
488+
// by userspace.
489+
throw entangledActionThenable;
490+
}
491+
}
492+
}
493+
466494
export function processUpdateQueue<State>(
467495
workInProgress: Fiber,
468496
props: any,
469497
instance: any,
470498
renderLanes: Lanes,
471499
): void {
500+
didReadFromEntangledAsyncAction = false;
501+
472502
// This is always non-null on a ClassComponent or HostRoot
473503
const queue: UpdateQueue<State> = (workInProgress.updateQueue: any);
474504

@@ -571,6 +601,13 @@ export function processUpdateQueue<State>(
571601
} else {
572602
// This update does have sufficient priority.
573603

604+
// Check if this update is part of a pending async action. If so,
605+
// we'll need to suspend until the action has finished, so that it's
606+
// batched together with future updates in the same action.
607+
if (updateLane !== NoLane && updateLane === peekEntangledActionLane()) {
608+
didReadFromEntangledAsyncAction = true;
609+
}
610+
574611
if (newLastBaseUpdate !== null) {
575612
const clone: Update<State> = {
576613
// This update is going to be committed so we never want uncommit

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

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1461,4 +1461,184 @@ describe('ReactAsyncActions', () => {
14611461
expect(root).toMatchRenderedOutput(<span>C</span>);
14621462
},
14631463
);
1464+
1465+
// @gate enableAsyncActions
1466+
test(
1467+
'updates in an async action are entangled even if useTransition hook ' +
1468+
'is unmounted before it finishes (class component)',
1469+
async () => {
1470+
let startTransition;
1471+
function Updater() {
1472+
const [isPending, _start] = useTransition();
1473+
startTransition = _start;
1474+
return (
1475+
<span>
1476+
<Text text={'Pending: ' + isPending} />
1477+
</span>
1478+
);
1479+
}
1480+
1481+
let setText;
1482+
class Sibling extends React.Component {
1483+
state = {text: 'A'};
1484+
render() {
1485+
setText = text => this.setState({text});
1486+
return (
1487+
<span>
1488+
<Text text={this.state.text} />
1489+
</span>
1490+
);
1491+
}
1492+
}
1493+
1494+
function App({showUpdater}) {
1495+
return (
1496+
<>
1497+
{showUpdater ? <Updater /> : null}
1498+
<Sibling />
1499+
</>
1500+
);
1501+
}
1502+
1503+
const root = ReactNoop.createRoot();
1504+
await act(() => {
1505+
root.render(<App showUpdater={true} />);
1506+
});
1507+
assertLog(['Pending: false', 'A']);
1508+
expect(root).toMatchRenderedOutput(
1509+
<>
1510+
<span>Pending: false</span>
1511+
<span>A</span>
1512+
</>,
1513+
);
1514+
1515+
// Start an async action that has multiple updates with async
1516+
// operations in between.
1517+
await act(() => {
1518+
startTransition(async () => {
1519+
Scheduler.log('Async action started');
1520+
startTransition(() => setText('B'));
1521+
1522+
await getText('Wait before updating to C');
1523+
1524+
Scheduler.log('Async action ended');
1525+
startTransition(() => setText('C'));
1526+
});
1527+
});
1528+
assertLog(['Async action started', 'Pending: true']);
1529+
expect(root).toMatchRenderedOutput(
1530+
<>
1531+
<span>Pending: true</span>
1532+
<span>A</span>
1533+
</>,
1534+
);
1535+
1536+
// Delete the component that contains the useTransition hook. This
1537+
// component no longer blocks the transition from completing. But the
1538+
// pending update to Sibling should not be allowed to finish, because it's
1539+
// part of the async action.
1540+
await act(() => {
1541+
root.render(<App showUpdater={false} />);
1542+
});
1543+
assertLog(['A']);
1544+
expect(root).toMatchRenderedOutput(<span>A</span>);
1545+
1546+
// Finish the async action. Notice the intermediate B state was never
1547+
// shown, because it was batched with the update that came later in the
1548+
// same action.
1549+
await act(() => resolveText('Wait before updating to C'));
1550+
assertLog(['Async action ended', 'C']);
1551+
expect(root).toMatchRenderedOutput(<span>C</span>);
1552+
1553+
// Check that subsequent updates are unaffected.
1554+
await act(() => setText('D'));
1555+
assertLog(['D']);
1556+
expect(root).toMatchRenderedOutput(<span>D</span>);
1557+
},
1558+
);
1559+
1560+
// @gate enableAsyncActions
1561+
test(
1562+
'updates in an async action are entangled even if useTransition hook ' +
1563+
'is unmounted before it finishes (root update)',
1564+
async () => {
1565+
let startTransition;
1566+
function Updater() {
1567+
const [isPending, _start] = useTransition();
1568+
startTransition = _start;
1569+
return (
1570+
<span>
1571+
<Text text={'Pending: ' + isPending} />
1572+
</span>
1573+
);
1574+
}
1575+
1576+
let setShowUpdater;
1577+
function App({text}) {
1578+
const [showUpdater, _setShowUpdater] = useState(true);
1579+
setShowUpdater = _setShowUpdater;
1580+
return (
1581+
<>
1582+
{showUpdater ? <Updater /> : null}
1583+
<span>
1584+
<Text text={text} />
1585+
</span>
1586+
</>
1587+
);
1588+
}
1589+
1590+
const root = ReactNoop.createRoot();
1591+
await act(() => {
1592+
root.render(<App text="A" />);
1593+
});
1594+
assertLog(['Pending: false', 'A']);
1595+
expect(root).toMatchRenderedOutput(
1596+
<>
1597+
<span>Pending: false</span>
1598+
<span>A</span>
1599+
</>,
1600+
);
1601+
1602+
// Start an async action that has multiple updates with async
1603+
// operations in between.
1604+
await act(() => {
1605+
startTransition(async () => {
1606+
Scheduler.log('Async action started');
1607+
startTransition(() => root.render(<App text="B" />));
1608+
1609+
await getText('Wait before updating to C');
1610+
1611+
Scheduler.log('Async action ended');
1612+
startTransition(() => root.render(<App text="C" />));
1613+
});
1614+
});
1615+
assertLog(['Async action started', 'Pending: true']);
1616+
expect(root).toMatchRenderedOutput(
1617+
<>
1618+
<span>Pending: true</span>
1619+
<span>A</span>
1620+
</>,
1621+
);
1622+
1623+
// Delete the component that contains the useTransition hook. This
1624+
// component no longer blocks the transition from completing. But the
1625+
// pending update to Sibling should not be allowed to finish, because it's
1626+
// part of the async action.
1627+
await act(() => setShowUpdater(false));
1628+
assertLog(['A']);
1629+
expect(root).toMatchRenderedOutput(<span>A</span>);
1630+
1631+
// Finish the async action. Notice the intermediate B state was never
1632+
// shown, because it was batched with the update that came later in the
1633+
// same action.
1634+
await act(() => resolveText('Wait before updating to C'));
1635+
assertLog(['Async action ended', 'C']);
1636+
expect(root).toMatchRenderedOutput(<span>C</span>);
1637+
1638+
// Check that subsequent updates are unaffected.
1639+
await act(() => root.render(<App text="D" />));
1640+
assertLog(['D']);
1641+
expect(root).toMatchRenderedOutput(<span>D</span>);
1642+
},
1643+
);
14641644
});

0 commit comments

Comments
 (0)