Skip to content

Commit

Permalink
Batch async actions even if useTransition is unmounted
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 committed Jan 24, 2024
1 parent 42a0d5b commit 0ffbed7
Show file tree
Hide file tree
Showing 3 changed files with 297 additions and 137 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;
}
56 changes: 36 additions & 20 deletions packages/react-reconciler/src/ReactFiberHooks.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,9 +145,10 @@ import {
import type {ThenableState} from './ReactFiberThenable';
import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';
import {
requestAsyncActionContext,
requestSyncActionContext,
entangleAsyncAction,
peekEntangledActionLane,
peekEntangledActionThenable,
chainThenableValue,
} from './ReactFiberAsyncAction';
import {HostTransitionContext} from './ReactFiberHostContext';
import {requestTransitionLane} from './ReactFiberRootScheduler';
Expand Down Expand Up @@ -1274,6 +1275,7 @@ function updateReducerImpl<S, A>(
let newBaseQueueFirst = null;
let newBaseQueueLast: Update<S, A> | null = null;
let update = first;
let didReadFromEntangledAsyncAction = false;
do {
// An extra OffscreenLane bit is added to updates that were made to
// a hidden tree, so that we can distinguish them from updates that were
Expand Down Expand Up @@ -1317,6 +1319,13 @@ function updateReducerImpl<S, A>(
} 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;
}

// Check if this is an optimistic update.
const revertLane = update.revertLane;
if (!enableAsyncActions || revertLane === NoLane) {
Expand Down Expand Up @@ -1407,6 +1416,22 @@ function updateReducerImpl<S, A>(
// different from the current state.
if (!is(newState, hook.memoizedState)) {
markWorkInProgressReceivedUpdate();

// 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;
}
}
}

hook.memoizedState = newState;
Expand Down Expand Up @@ -1964,13 +1989,10 @@ function runFormStateAction<S, P>(
() => finishRunningFormStateAction(actionQueue, (setState: any)),
);

const entangledResult = requestAsyncActionContext<S>(thenable, null);
setState((entangledResult: any));
entangleAsyncAction<Awaited<S>>(thenable);
setState((thenable: any));
} else {
// This is either `returnValue` or a thenable that resolves to
// `returnValue`, depending on whether we're inside an async action scope.
const entangledResult = requestSyncActionContext<S>(returnValue, null);
setState((entangledResult: any));
setState((returnValue: any));

const nextState = ((returnValue: any): Awaited<S>);
actionQueue.state = nextState;
Expand Down Expand Up @@ -2832,22 +2854,16 @@ function startTransition<S>(
typeof returnValue.then === 'function'
) {
const thenable = ((returnValue: any): Thenable<mixed>);
// This is a thenable that resolves to `finishedState` once the async
// action scope has finished.
const entangledResult = requestAsyncActionContext(
entangleAsyncAction<mixed>(thenable);
// Create a thenable that resolves to `finishedState` once the async
// action has completed.
const thenableForFinishedState = chainThenableValue(
thenable,
finishedState,
);
dispatchSetState(fiber, queue, entangledResult);
dispatchSetState(fiber, queue, (thenableForFinishedState: any));
} else {
// This is either `finishedState` or a thenable that resolves to
// `finishedState`, depending on whether we're inside an async
// action scope.
const entangledResult = requestSyncActionContext(
returnValue,
finishedState,
);
dispatchSetState(fiber, queue, entangledResult);
dispatchSetState(fiber, queue, finishedState);
}
} else {
// Async actions are not enabled.
Expand Down
Loading

0 comments on commit 0ffbed7

Please sign in to comment.