Initial (client-only) async actions support

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

@@ -15,6 +15,7 @@ import type {
   StartTransitionOptions,
   Usable,
   Thenable,
+  RejectedThenable,
 } from 'shared/ReactTypes';
 import type {
   Fiber,
@@ -41,6 +42,7 @@ import {
   enableUseEffectEventHook,
   enableLegacyCache,
   debugRenderPhaseSideEffectsForStrictMode,
+  enableAsyncActions,
 } from 'shared/ReactFeatureFlags';
 import {
   REACT_CONTEXT_TYPE,
@@ -143,6 +145,7 @@ import {
 } from './ReactFiberThenable';
 import type {ThenableState} from './ReactFiberThenable';
 import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';
+import {requestAsyncActionContext} from './ReactFiberAsyncAction';
 
 const {ReactCurrentDispatcher, ReactCurrentBatchConfig} = ReactSharedInternals;
 
@@ -947,38 +950,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 +2405,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 +2432,36 @@ function startTransition(
   }
 
   try {
-    setPending(false);
-    callback();
+    if (enableAsyncActions) {
+      const returnValue = callback();
+
+      // `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);
+      callback();
+    }
+  } catch (error) {
+    if (enableAsyncActions) {
+      // This is a trick to get the `useTransition` hook to rethrow the error.
+      // When it unwraps the thenable with the `use` algorithm, the error
+      // will be thrown.
+      const rejectedThenable: RejectedThenable<boolean> = {
+        then() {},
+        status: 'rejected',
+        reason: error,
+      };
+      setPending(rejectedThenable);
+    } else {
+      // The error rethrowing behavior is only enabled when the async actions
+      // feature is on, even for sync actions.
+      throw error;
+    }
   } finally {
     setCurrentUpdatePriority(previousPriority);
 
@@ -2454,31 +2487,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;
-}