Skip to content

Commit 11c9fd0

Browse files
acdliteacdlite
authored
Batch async actions even if useTransition is unmounted (#28078)
If there are multiple updates inside an async action, they should all be rendered in the same batch, even if they are separate by an async operation (`await`). We currently implement this by suspending in the `useTransition` hook to block the update from committing until all possible updates have been scheduled by the action. The reason we did it this way is so you can "cancel" an action by navigating away from the UI that triggered it. The problem with that approach, though, is that even if you navigate away from the `useTransition` hook, the action may have updated shared parts of the UI that are still in the tree. So we may need to continue suspending even after the `useTransition` hook is deleted. In other words, the lifetime of an async action scope is longer than the lifetime of a particular `useTransition` hook. The solution is to suspend whenever _any_ update that is part of the async action scope is unwrapped during render. So, inside useState and useReducer. This fixes a related issue where an optimistic update is reverted before the async action has finished, because we were relying on the `useTransition` hook to prevent the optimistic update from finishing. This also prepares us to support async actions being passed to the non-hook form of `startTransition` (though this isn't implemented yet).
1 parent 2efb142 commit 11c9fd0

8 files changed

+582
-173
lines changed

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

+70-117
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99

1010
import type {
1111
Thenable,
12-
PendingThenable,
1312
FulfilledThenable,
1413
RejectedThenable,
1514
} from 'shared/ReactTypes';
@@ -32,111 +31,32 @@ let currentEntangledListeners: Array<() => mixed> | null = null;
3231
let currentEntangledPendingCount: number = 0;
3332
// The transition lane shared by all updates in the entangled scope.
3433
let currentEntangledLane: Lane = NoLane;
34+
// A thenable that resolves when the entangled scope completes. It does not
35+
// resolve to a particular value because it's only used for suspending the UI
36+
// until the async action scope has completed.
37+
let currentEntangledActionThenable: Thenable<void> | null = null;
3538

36-
export function requestAsyncActionContext<S>(
37-
actionReturnValue: Thenable<any>,
38-
// If this is provided, this resulting thenable resolves to this value instead
39-
// of the return value of the action. This is a perf trick to avoid composing
40-
// an extra async function.
41-
overrideReturnValue: S | null,
42-
): Thenable<S> {
43-
// This is an async action.
44-
//
45-
// Return a thenable that resolves once the action scope (i.e. the async
46-
// function passed to startTransition) has finished running.
47-
48-
const thenable: Thenable<S> = (actionReturnValue: any);
49-
let entangledListeners;
39+
export function entangleAsyncAction<S>(thenable: Thenable<S>): Thenable<S> {
40+
// `thenable` is the return value of the async action scope function. Create
41+
// a combined thenable that resolves once every entangled scope function
42+
// has finished.
5043
if (currentEntangledListeners === null) {
5144
// There's no outer async action scope. Create a new one.
52-
entangledListeners = currentEntangledListeners = [];
45+
const entangledListeners = (currentEntangledListeners = []);
5346
currentEntangledPendingCount = 0;
5447
currentEntangledLane = requestTransitionLane();
55-
} else {
56-
entangledListeners = currentEntangledListeners;
48+
const entangledThenable: Thenable<void> = {
49+
status: 'pending',
50+
value: undefined,
51+
then(resolve: void => mixed) {
52+
entangledListeners.push(resolve);
53+
},
54+
};
55+
currentEntangledActionThenable = entangledThenable;
5756
}
58-
5957
currentEntangledPendingCount++;
60-
61-
// Create a thenable that represents the result of this action, but doesn't
62-
// resolve until the entire entangled scope has finished.
63-
//
64-
// Expressed using promises:
65-
// const [thisResult] = await Promise.all([thisAction, entangledAction]);
66-
// return thisResult;
67-
const resultThenable = createResultThenable<S>(entangledListeners);
68-
69-
let resultStatus = 'pending';
70-
let resultValue;
71-
let rejectedReason;
72-
thenable.then(
73-
(value: S) => {
74-
resultStatus = 'fulfilled';
75-
resultValue = overrideReturnValue !== null ? overrideReturnValue : value;
76-
pingEngtangledActionScope();
77-
},
78-
error => {
79-
resultStatus = 'rejected';
80-
rejectedReason = error;
81-
pingEngtangledActionScope();
82-
},
83-
);
84-
85-
// Attach a listener to fill in the result.
86-
entangledListeners.push(() => {
87-
switch (resultStatus) {
88-
case 'fulfilled': {
89-
const fulfilledThenable: FulfilledThenable<S> = (resultThenable: any);
90-
fulfilledThenable.status = 'fulfilled';
91-
fulfilledThenable.value = resultValue;
92-
break;
93-
}
94-
case 'rejected': {
95-
const rejectedThenable: RejectedThenable<S> = (resultThenable: any);
96-
rejectedThenable.status = 'rejected';
97-
rejectedThenable.reason = rejectedReason;
98-
break;
99-
}
100-
case 'pending':
101-
default: {
102-
// The listener above should have been called first, so `resultStatus`
103-
// should already be set to the correct value.
104-
throw new Error(
105-
'Thenable should have already resolved. This ' + 'is a bug in React.',
106-
);
107-
}
108-
}
109-
});
110-
111-
return resultThenable;
112-
}
113-
114-
export function requestSyncActionContext<S>(
115-
actionReturnValue: any,
116-
// If this is provided, this resulting thenable resolves to this value instead
117-
// of the return value of the action. This is a perf trick to avoid composing
118-
// an extra async function.
119-
overrideReturnValue: S | null,
120-
): Thenable<S> | S {
121-
const resultValue: S =
122-
overrideReturnValue !== null
123-
? overrideReturnValue
124-
: (actionReturnValue: any);
125-
// This is not an async action, but it may be part of an outer async action.
126-
if (currentEntangledListeners === null) {
127-
return resultValue;
128-
} else {
129-
// Return a thenable that does not resolve until the entangled actions
130-
// have finished.
131-
const entangledListeners = currentEntangledListeners;
132-
const resultThenable = createResultThenable<S>(entangledListeners);
133-
entangledListeners.push(() => {
134-
const fulfilledThenable: FulfilledThenable<S> = (resultThenable: any);
135-
fulfilledThenable.status = 'fulfilled';
136-
fulfilledThenable.value = resultValue;
137-
});
138-
return resultThenable;
139-
}
58+
thenable.then(pingEngtangledActionScope, pingEngtangledActionScope);
59+
return thenable;
14060
}
14161

14262
function pingEngtangledActionScope() {
@@ -146,41 +66,74 @@ function pingEngtangledActionScope() {
14666
) {
14767
// All the actions have finished. Close the entangled async action scope
14868
// and notify all the listeners.
69+
if (currentEntangledActionThenable !== null) {
70+
const fulfilledThenable: FulfilledThenable<void> =
71+
(currentEntangledActionThenable: any);
72+
fulfilledThenable.status = 'fulfilled';
73+
}
14974
const listeners = currentEntangledListeners;
15075
currentEntangledListeners = null;
15176
currentEntangledLane = NoLane;
77+
currentEntangledActionThenable = null;
15278
for (let i = 0; i < listeners.length; i++) {
15379
const listener = listeners[i];
15480
listener();
15581
}
15682
}
15783
}
15884

159-
function createResultThenable<S>(
160-
entangledListeners: Array<() => mixed>,
161-
): Thenable<S> {
162-
// Waits for the entangled async action to complete, then resolves to the
163-
// result of an individual action.
164-
const resultThenable: PendingThenable<S> = {
85+
export function chainThenableValue<T>(
86+
thenable: Thenable<T>,
87+
result: T,
88+
): Thenable<T> {
89+
// Equivalent to: Promise.resolve(thenable).then(() => result), except we can
90+
// cheat a bit since we know that that this thenable is only ever consumed
91+
// by React.
92+
//
93+
// We don't technically require promise support on the client yet, hence this
94+
// extra code.
95+
const listeners = [];
96+
const thenableWithOverride: Thenable<T> = {
16597
status: 'pending',
16698
value: null,
16799
reason: null,
168-
then(resolve: S => mixed) {
169-
// This is a bit of a cheat. `resolve` expects a value of type `S` to be
170-
// passed, but because we're instrumenting the `status` field ourselves,
171-
// and we know this thenable will only be used by React, we also know
172-
// the value isn't actually needed. So we add the resolve function
173-
// directly to the entangled listeners.
174-
//
175-
// This is also why we don't need to check if the thenable is still
176-
// pending; the Suspense implementation already performs that check.
177-
const ping: () => mixed = (resolve: any);
178-
entangledListeners.push(ping);
100+
then(resolve: T => mixed) {
101+
listeners.push(resolve);
179102
},
180103
};
181-
return resultThenable;
104+
thenable.then(
105+
(value: T) => {
106+
const fulfilledThenable: FulfilledThenable<T> =
107+
(thenableWithOverride: any);
108+
fulfilledThenable.status = 'fulfilled';
109+
fulfilledThenable.value = result;
110+
for (let i = 0; i < listeners.length; i++) {
111+
const listener = listeners[i];
112+
listener(result);
113+
}
114+
},
115+
error => {
116+
const rejectedThenable: RejectedThenable<T> = (thenableWithOverride: any);
117+
rejectedThenable.status = 'rejected';
118+
rejectedThenable.reason = error;
119+
for (let i = 0; i < listeners.length; i++) {
120+
const listener = listeners[i];
121+
// This is a perf hack where we call the `onFulfill` ping function
122+
// instead of `onReject`, because we know that React is the only
123+
// consumer of these promises, and it passes the same listener to both.
124+
// We also know that it will read the error directly off the
125+
// `.reason` field.
126+
listener((undefined: any));
127+
}
128+
},
129+
);
130+
return thenableWithOverride;
182131
}
183132

184133
export function peekEntangledActionLane(): Lane {
185134
return currentEntangledLane;
186135
}
136+
137+
export function peekEntangledActionThenable(): Thenable<void> | null {
138+
return currentEntangledActionThenable;
139+
}

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

+7
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

+4
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

+37
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

0 commit comments

Comments
 (0)