Initial (client-only) async actions support

Implements initial (client-only) support for async actions behind a
flag. This is an experimental feature and the design isn't completely
finalized but we're getting closer. It will be layered alongside other
features we're working on, so it may not feel complete when considered
in isolation.

The basic description is you can pass an async function to
`startTransition` and all the transition updates that are scheduled
inside that async function will be grouped together. The `isPending`
flag will be set to true immediately, and only set back to false once
the async action has completed (as well as all the updates that
it triggers).

The ideal behavior would be that all updates spawned by the async
action are automatically inferred and grouped together; however, doing
this properly requires the upcoming (stage 2) Async Context API, which
is not yet implemented by browsers. In the meantime, we will fake this
by grouping together all transition updates that occur until the async
function has terminated. This can lead to overgrouping between unrelated
actions, which is not wrong per se, just not ideal.

If the `useTransition` hook is removed from the UI before an async
action has completed — for example, if the user navigates to a new page
— subsequent transitions will no longer be grouped with together with
that action.

Another consequence of the lack of Async Context is that if you call
`setState` inside an action but after an `await`, it must be wrapped in
`startTransition` in order to be grouped properly. If we didn't require
this, then there would be no way to distinguish action updates from
urgent updates caused by user input, too. This is an unfortunate footgun
but we can likely detect the most common mistakes using a lint rule.

Once Async Context lands in browsers, we can start warning in dev if we
detect an update that hasn't been wrapped in `startTransition`. Then,
longer term, once the feature is ubiquitous, we can rely on it for real
and allow you to call `setState` without the additional wrapper.

Things that are _not_ yet implemented in this commit, but will be added
as follow ups:

- Support for non-hook from of `startTransition`
- Canceling the async action scope if the `useTransition` hook
  is deleted from the UI
- Anything related to server actions

Author: acdlite

packages/react-reconciler/src/ReactFiberAsyncAction.js

@@ -0,0 +1,130 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {Wakeable} from 'shared/ReactTypes';
+import type {Lane} from './ReactFiberLane';
+import {requestTransitionLane} from './ReactFiberRootScheduler';
+
+interface AsyncActionImpl {
+  lane: Lane;
+  listeners: Array<(false) => mixed>;
+  count: number;
+  then(
+    onFulfill: (value: boolean) => mixed,
+    onReject: (error: mixed) => mixed,
+  ): void;
+}
+
+interface PendingAsyncAction extends AsyncActionImpl {
+  status: 'pending';
+}
+
+interface FulfilledAsyncAction extends AsyncActionImpl {
+  status: 'fulfilled';
+  value: boolean;
+}
+
+interface RejectedAsyncAction extends AsyncActionImpl {
+  status: 'rejected';
+  reason: mixed;
+}
+
+type AsyncAction =
+  | PendingAsyncAction
+  | FulfilledAsyncAction
+  | RejectedAsyncAction;
+
+let currentAsyncAction: AsyncAction | null = null;
+
+export function requestAsyncActionContext(
+  actionReturnValue: mixed,
+): AsyncAction | false {
+  if (
+    actionReturnValue !== null &&
+    typeof actionReturnValue === 'object' &&
+    typeof actionReturnValue.then === 'function'
+  ) {
+    // 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. The fulfilled
+    // value is `false` to represent that the action is not pending.
+    const thenable: Wakeable = (actionReturnValue: any);
+    if (currentAsyncAction === null) {
+      // There's no outer async action scope. Create a new one.
+      const asyncAction: AsyncAction = {
+        lane: requestTransitionLane(),
+        listeners: [],
+        count: 0,
+        status: 'pending',
+        value: false,
+        reason: undefined,
+        then(resolve: boolean => mixed) {
+          asyncAction.listeners.push(resolve);
+        },
+      };
+      attachPingListeners(thenable, asyncAction);
+      currentAsyncAction = asyncAction;
+      return asyncAction;
+    } else {
+      // Inherit the outer scope.
+      const asyncAction: AsyncAction = (currentAsyncAction: any);
+      attachPingListeners(thenable, asyncAction);
+      return asyncAction;
+    }
+  } else {
+    // This is not an async action, but it may be part of an outer async action.
+    if (currentAsyncAction === null) {
+      // There's no outer async action scope.
+      return false;
+    } else {
+      // Inherit the outer scope.
+      return currentAsyncAction;
+    }
+  }
+}
+
+export function peekAsyncActionContext(): AsyncAction | null {
+  return currentAsyncAction;
+}
+
+function attachPingListeners(thenable: Wakeable, asyncAction: AsyncAction) {
+  asyncAction.count++;
+  thenable.then(
+    () => {
+      if (--asyncAction.count === 0) {
+        const fulfilledAsyncAction: FulfilledAsyncAction = (asyncAction: any);
+        fulfilledAsyncAction.status = 'fulfilled';
+        completeAsyncActionScope(asyncAction);
+      }
+    },
+    (error: mixed) => {
+      if (--asyncAction.count === 0) {
+        const rejectedAsyncAction: RejectedAsyncAction = (asyncAction: any);
+        rejectedAsyncAction.status = 'rejected';
+        rejectedAsyncAction.reason = error;
+        completeAsyncActionScope(asyncAction);
+      }
+    },
+  );
+  return asyncAction;
+}
+
+function completeAsyncActionScope(action: AsyncAction) {
+  if (currentAsyncAction === action) {
+    currentAsyncAction = null;
+  }
+
+  const listeners = action.listeners;
+  action.listeners = [];
+  for (let i = 0; i < listeners.length; i++) {
+    const listener = listeners[i];
+    listener(false);
+  }
+}

packages/react-reconciler/src/ReactFiberHooks.js

@@ -41,6 +41,7 @@ import {
   enableUseEffectEventHook,
   enableLegacyCache,
   debugRenderPhaseSideEffectsForStrictMode,
+  enableAsyncActions,
 } from 'shared/ReactFeatureFlags';
 import {
   REACT_CONTEXT_TYPE,
@@ -143,6 +144,7 @@ import {
 } from './ReactFiberThenable';
 import type {ThenableState} from './ReactFiberThenable';
 import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';
+import {requestAsyncActionContext} from './ReactFiberAsyncAction';
 
 const {ReactCurrentDispatcher, ReactCurrentBatchConfig} = ReactSharedInternals;
 
@@ -947,38 +949,40 @@ if (enableUseMemoCacheHook) {
   };
 }
 
+function useThenable<T>(thenable: Thenable<T>): T {
+  // Track the position of the thenable within this fiber.
+  const index = thenableIndexCounter;
+  thenableIndexCounter += 1;
+  if (thenableState === null) {
+    thenableState = createThenableState();
+  }
+  const result = trackUsedThenable(thenableState, thenable, index);
+  if (
+    currentlyRenderingFiber.alternate === null &&
+    (workInProgressHook === null
+      ? currentlyRenderingFiber.memoizedState === null
+      : workInProgressHook.next === null)
+  ) {
+    // Initial render, and either this is the first time the component is
+    // called, or there were no Hooks called after this use() the previous
+    // time (perhaps because it threw). Subsequent Hook calls should use the
+    // mount dispatcher.
+    if (__DEV__) {
+      ReactCurrentDispatcher.current = HooksDispatcherOnMountInDEV;
+    } else {
+      ReactCurrentDispatcher.current = HooksDispatcherOnMount;
+    }
+  }
+  return result;
+}
+
 function use<T>(usable: Usable<T>): T {
   if (usable !== null && typeof usable === 'object') {
     // $FlowFixMe[method-unbinding]
     if (typeof usable.then === 'function') {
       // This is a thenable.
       const thenable: Thenable<T> = (usable: any);
-
-      // Track the position of the thenable within this fiber.
-      const index = thenableIndexCounter;
-      thenableIndexCounter += 1;
-
-      if (thenableState === null) {
-        thenableState = createThenableState();
-      }
-      const result = trackUsedThenable(thenableState, thenable, index);
-      if (
-        currentlyRenderingFiber.alternate === null &&
-        (workInProgressHook === null
-          ? currentlyRenderingFiber.memoizedState === null
-          : workInProgressHook.next === null)
-      ) {
-        // Initial render, and either this is the first time the component is
-        // called, or there were no Hooks called after this use() the previous
-        // time (perhaps because it threw). Subsequent Hook calls should use the
-        // mount dispatcher.
-        if (__DEV__) {
-          ReactCurrentDispatcher.current = HooksDispatcherOnMountInDEV;
-        } else {
-          ReactCurrentDispatcher.current = HooksDispatcherOnMount;
-        }
-      }
-      return result;
+      return useThenable(thenable);
     } else if (
       usable.$$typeof === REACT_CONTEXT_TYPE ||
       usable.$$typeof === REACT_SERVER_CONTEXT_TYPE
@@ -2400,8 +2404,8 @@ function updateDeferredValueImpl<T>(hook: Hook, prevValue: T, value: T): T {
 }
 
 function startTransition(
-  setPending: boolean => void,
-  callback: () => void,
+  setPending: (Thenable<boolean> | boolean) => void,
+  callback: () => mixed,
   options?: StartTransitionOptions,
 ): void {
   const previousPriority = getCurrentUpdatePriority();
@@ -2427,8 +2431,18 @@ function startTransition(
   }
 
   try {
-    setPending(false);
-    callback();
+    const returnValue = callback();
+    if (enableAsyncActions) {
+      // `isPending` is either `false` or a thenable that resolves to `false`,
+      // depending on whether the action scope is an async function. In the
+      // async case, the resulting render will suspend until the async action
+      // scope has finished.
+      const isPending = requestAsyncActionContext(returnValue);
+      setPending(isPending);
+    } else {
+      // Async actions are not enabled.
+      setPending(false);
+    }
   } finally {
     setCurrentUpdatePriority(previousPriority);
 
@@ -2454,31 +2468,41 @@ function mountTransition(): [
   boolean,
   (callback: () => void, options?: StartTransitionOptions) => void,
 ] {
-  const [isPending, setPending] = mountState(false);
+  const [, setPending] = mountState((false: Thenable<boolean> | boolean));
   // The `start` method never changes.
   const start = startTransition.bind(null, setPending);
   const hook = mountWorkInProgressHook();
   hook.memoizedState = start;
-  return [isPending, start];
+  return [false, start];
 }
 
 function updateTransition(): [
   boolean,
   (callback: () => void, options?: StartTransitionOptions) => void,
 ] {
-  const [isPending] = updateState(false);
+  const [booleanOrThenable] = updateState(false);
   const hook = updateWorkInProgressHook();
   const start = hook.memoizedState;
+  const isPending =
+    typeof booleanOrThenable === 'boolean'
+      ? booleanOrThenable
+      : // This will suspend until the async action scope has finished.
+        useThenable(booleanOrThenable);
   return [isPending, start];
 }
 
 function rerenderTransition(): [
   boolean,
   (callback: () => void, options?: StartTransitionOptions) => void,
 ] {
-  const [isPending] = rerenderState(false);
+  const [booleanOrThenable] = rerenderState(false);
   const hook = updateWorkInProgressHook();
   const start = hook.memoizedState;
+  const isPending =
+    typeof booleanOrThenable === 'boolean'
+      ? booleanOrThenable
+      : // This will suspend until the async action scope has finished.
+        useThenable(booleanOrThenable);
   return [isPending, start];
 }
 

packages/react-reconciler/src/ReactFiberRootScheduler.js

@@ -22,6 +22,7 @@ import {
   markStarvedLanesAsExpired,
   markRootEntangled,
   mergeLanes,
+  claimNextTransitionLane,
 } from './ReactFiberLane';
 import {
   CommitContext,
@@ -78,7 +79,7 @@ let mightHavePendingSyncWork: boolean = false;
 
 let isFlushingWork: boolean = false;
 
-let currentEventTransitionLane: Lane = NoLanes;
+let currentEventTransitionLane: Lane = NoLane;
 
 export function ensureRootIsScheduled(root: FiberRoot): void {
   // This function is called whenever a root receives an update. It does two
@@ -491,10 +492,17 @@ function scheduleImmediateTask(cb: () => mixed) {
   }
 }
 
-export function getCurrentEventTransitionLane(): Lane {
+export function requestTransitionLane(): Lane {
+  // The algorithm for assigning an update to a lane should be stable for all
+  // updates at the same priority within the same event. To do this, the
+  // inputs to the algorithm must be the same.
+  //
+  // The trick we use is to cache the first of each of these inputs within an
+  // event. Then reset the cached values once we can be sure the event is
+  // over. Our heuristic for that is whenever we enter a concurrent work loop.
+  if (currentEventTransitionLane === NoLane) {
+    // All transitions within the same event are assigned the same lane.
+    currentEventTransitionLane = claimNextTransitionLane();
+  }
   return currentEventTransitionLane;
 }
-
-export function setCurrentEventTransitionLane(lane: Lane): void {
-  currentEventTransitionLane = lane;
-}

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -129,7 +129,6 @@ import {
   NoLanes,
   NoLane,
   SyncLane,
-  claimNextTransitionLane,
   claimNextRetryLane,
   includesSyncLane,
   isSubsetOfLanes,
@@ -278,10 +277,10 @@ import {
   flushSyncWorkOnAllRoots,
   flushSyncWorkOnLegacyRootsOnly,
   getContinuationForRoot,
-  getCurrentEventTransitionLane,
-  setCurrentEventTransitionLane,
+  requestTransitionLane,
 } from './ReactFiberRootScheduler';
 import {getMaskedContext, getUnmaskedContext} from './ReactFiberContext';
+import {peekAsyncActionContext} from './ReactFiberAsyncAction';
 
 const PossiblyWeakMap = typeof WeakMap === 'function' ? WeakMap : Map;
 
@@ -633,18 +632,15 @@ export function requestUpdateLane(fiber: Fiber): Lane {
 
       transition._updatedFibers.add(fiber);
     }
-    // The algorithm for assigning an update to a lane should be stable for all
-    // updates at the same priority within the same event. To do this, the
-    // inputs to the algorithm must be the same.
-    //
-    // The trick we use is to cache the first of each of these inputs within an
-    // event. Then reset the cached values once we can be sure the event is
-    // over. Our heuristic for that is whenever we enter a concurrent work loop.
-    if (getCurrentEventTransitionLane() === NoLane) {
-      // All transitions within the same event are assigned the same lane.
-      setCurrentEventTransitionLane(claimNextTransitionLane());
-    }
-    return getCurrentEventTransitionLane();
+
+    const asyncAction = peekAsyncActionContext();
+    return asyncAction !== null
+      ? // We're inside an async action scope. Reuse the same lane.
+        asyncAction.lane
+      : // We may or may not be inside an async action scope. If we are, this
+        // is the first update in that scope. Either way, we need to get a
+        // fresh transition lane.
+        requestTransitionLane();
   }
 
   // Updates originating inside certain React methods, like flushSync, have