Skip to content

Commit

Permalink
Batch async actions even if useTransition is unmounted (#28078)
Browse files Browse the repository at this point in the history 
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).
  • Loading branch information
@acdlite
acdlite authored Jan 25, 2024
1 parent 2efb142 commit 11c9fd0
Show file tree
Hide file tree
Showing 8 changed files with 582 additions and 173 deletions.
187 changes: 70 additions & 117 deletions packages/react-reconciler/src/ReactFiberAsyncAction.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@

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

export function requestAsyncActionContext<S>(
actionReturnValue: Thenable<any>,
// If this is provided, this resulting thenable resolves to this value instead
// of the return value of the action. This is a perf trick to avoid composing
// an extra async function.
overrideReturnValue: S | null,
): Thenable<S> {
// This is an async action.
//
// Return a thenable that resolves once the action scope (i.e. the async
// function passed to startTransition) has finished running.

const thenable: Thenable<S> = (actionReturnValue: any);
let entangledListeners;
export function entangleAsyncAction<S>(thenable: Thenable<S>): Thenable<S> {
// `thenable` is the return value of the async action scope function. Create
// a combined thenable that resolves once every entangled scope function
// has finished.
if (currentEntangledListeners === null) {
// There's no outer async action scope. Create a new one.
entangledListeners = currentEntangledListeners = [];
const entangledListeners = (currentEntangledListeners = []);
currentEntangledPendingCount = 0;
currentEntangledLane = requestTransitionLane();
} else {
entangledListeners = currentEntangledListeners;
const entangledThenable: Thenable<void> = {
status: 'pending',
value: undefined,
then(resolve: void => mixed) {
entangledListeners.push(resolve);
},
};
currentEntangledActionThenable = entangledThenable;
}

currentEntangledPendingCount++;

// Create a thenable that represents the result of this action, but doesn't
// resolve until the entire entangled scope has finished.
//
// Expressed using promises:
// const [thisResult] = await Promise.all([thisAction, entangledAction]);
// return thisResult;
const resultThenable = createResultThenable<S>(entangledListeners);

let resultStatus = 'pending';
let resultValue;
let rejectedReason;
thenable.then(
(value: S) => {
resultStatus = 'fulfilled';
resultValue = overrideReturnValue !== null ? overrideReturnValue : value;
pingEngtangledActionScope();
},
error => {
resultStatus = 'rejected';
rejectedReason = error;
pingEngtangledActionScope();
},
);

// Attach a listener to fill in the result.
entangledListeners.push(() => {
switch (resultStatus) {
case 'fulfilled': {
const fulfilledThenable: FulfilledThenable<S> = (resultThenable: any);
fulfilledThenable.status = 'fulfilled';
fulfilledThenable.value = resultValue;
break;
}
case 'rejected': {
const rejectedThenable: RejectedThenable<S> = (resultThenable: any);
rejectedThenable.status = 'rejected';
rejectedThenable.reason = rejectedReason;
break;
}
case 'pending':
default: {
// The listener above should have been called first, so `resultStatus`
// should already be set to the correct value.
throw new Error(
'Thenable should have already resolved. This ' + 'is a bug in React.',
);
}
}
});

return resultThenable;
}

export function requestSyncActionContext<S>(
actionReturnValue: any,
// If this is provided, this resulting thenable resolves to this value instead
// of the return value of the action. This is a perf trick to avoid composing
// an extra async function.
overrideReturnValue: S | null,
): Thenable<S> | S {
const resultValue: S =
overrideReturnValue !== null
? overrideReturnValue
: (actionReturnValue: any);
// This is not an async action, but it may be part of an outer async action.
if (currentEntangledListeners === null) {
return resultValue;
} else {
// Return a thenable that does not resolve until the entangled actions
// have finished.
const entangledListeners = currentEntangledListeners;
const resultThenable = createResultThenable<S>(entangledListeners);
entangledListeners.push(() => {
const fulfilledThenable: FulfilledThenable<S> = (resultThenable: any);
fulfilledThenable.status = 'fulfilled';
fulfilledThenable.value = resultValue;
});
return resultThenable;
}
thenable.then(pingEngtangledActionScope, pingEngtangledActionScope);
return thenable;
}

function pingEngtangledActionScope() {
Expand All @@ -146,41 +66,74 @@ function pingEngtangledActionScope() {
) {
// All the actions have finished. Close the entangled async action scope
// and notify all the listeners.
if (currentEntangledActionThenable !== null) {
const fulfilledThenable: FulfilledThenable<void> =
(currentEntangledActionThenable: any);
fulfilledThenable.status = 'fulfilled';
}
const listeners = currentEntangledListeners;
currentEntangledListeners = null;
currentEntangledLane = NoLane;
currentEntangledActionThenable = null;
for (let i = 0; i < listeners.length; i++) {
const listener = listeners[i];
listener();
}
}
}

function createResultThenable<S>(
entangledListeners: Array<() => mixed>,
): Thenable<S> {
// Waits for the entangled async action to complete, then resolves to the
// result of an individual action.
const resultThenable: PendingThenable<S> = {
export function chainThenableValue<T>(
thenable: Thenable<T>,
result: T,
): Thenable<T> {
// Equivalent to: Promise.resolve(thenable).then(() => result), except we can
// cheat a bit since we know that that this thenable is only ever consumed
// by React.
//
// We don't technically require promise support on the client yet, hence this
// extra code.
const listeners = [];
const thenableWithOverride: Thenable<T> = {
status: 'pending',
value: null,
reason: null,
then(resolve: S => mixed) {
// This is a bit of a cheat. `resolve` expects a value of type `S` to be
// passed, but because we're instrumenting the `status` field ourselves,
// and we know this thenable will only be used by React, we also know
// the value isn't actually needed. So we add the resolve function
// directly to the entangled listeners.
//
// This is also why we don't need to check if the thenable is still
// pending; the Suspense implementation already performs that check.
const ping: () => mixed = (resolve: any);
entangledListeners.push(ping);
then(resolve: T => mixed) {
listeners.push(resolve);
},
};
return resultThenable;
thenable.then(
(value: T) => {
const fulfilledThenable: FulfilledThenable<T> =
(thenableWithOverride: any);
fulfilledThenable.status = 'fulfilled';
fulfilledThenable.value = result;
for (let i = 0; i < listeners.length; i++) {
const listener = listeners[i];
listener(result);
}
},
error => {
const rejectedThenable: RejectedThenable<T> = (thenableWithOverride: any);
rejectedThenable.status = 'rejected';
rejectedThenable.reason = error;
for (let i = 0; i < listeners.length; i++) {
const listener = listeners[i];
// This is a perf hack where we call the `onFulfill` ping function
// instead of `onReject`, because we know that React is the only
// consumer of these promises, and it passes the same listener to both.
// We also know that it will read the error directly off the
// `.reason` field.
listener((undefined: any));
}
},
);
return thenableWithOverride;
}

export function peekEntangledActionLane(): Lane {
return currentEntangledLane;
}

export function peekEntangledActionThenable(): Thenable<void> | null {
return currentEntangledActionThenable;
}
7 changes: 7 additions & 0 deletions packages/react-reconciler/src/ReactFiberBeginWork.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ import {
cloneUpdateQueue,
initializeUpdateQueue,
enqueueCapturedUpdate,
suspendIfUpdateReadFromEntangledAsyncAction,
} from './ReactFiberClassUpdateQueue';
import {
NoLane,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
4 changes: 4 additions & 0 deletions packages/react-reconciler/src/ReactFiberClassComponent.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ import {
ForceUpdate,
initializeUpdateQueue,
cloneUpdateQueue,
suspendIfUpdateReadFromEntangledAsyncAction,
} from './ReactFiberClassUpdateQueue';
import {NoLanes} from './ReactFiberLane';
import {
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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 &&
Expand Down Expand Up @@ -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 (
Expand Down
37 changes: 37 additions & 0 deletions packages/react-reconciler/src/ReactFiberClassUpdateQueue.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,10 @@ import {
import {setIsStrictModeForDevtools} from './ReactFiberDevToolsHook';

import assign from 'shared/assign';
import {
peekEntangledActionLane,
peekEntangledActionThenable,
} from './ReactFiberAsyncAction';

export type Update<State> = {
lane: Lane,
Expand Down Expand Up @@ -463,12 +467,38 @@ function getStateFromUpdate<State>(
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<State>(
workInProgress: Fiber,
props: any,
instance: any,
renderLanes: Lanes,
): void {
didReadFromEntangledAsyncAction = false;

// This is always non-null on a ClassComponent or HostRoot
const queue: UpdateQueue<State> = (workInProgress.updateQueue: any);

Expand Down Expand Up @@ -571,6 +601,13 @@ export function processUpdateQueue<State>(
} 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<State> = {
// This update is going to be committed so we never want uncommit
Expand Down
Loading

0 comments on commit 11c9fd0

Please sign in to comment.