Repository: facebook/react
Pull request: useActionState: Transfer transition context

Mini-refactor of useActionState to only wrap the action in a transition context if the dispatch is called during a transition. Conceptually, the action starts as soon as the dispatch is called, even if the action is queued until earlier ones finish.

We will also warn if an async action is dispatched outside of a transition, since that is almost certainly a mistake. Ideally we would automatically upgrade these to a transition, but we don't have a great way to tell if the action is async until after it's already run.

Changed files:
filename: packages/react-dom/src/__tests__/ReactDOMForm-test.js
patch: @@ -1020,15 +1020,15 @@ describe('ReactDOMForm', () => {
     assertLog(['0']);
     expect(container.textContent).toBe('0');
 
-    await act(() => dispatch('increment'));
+    await act(() => startTransition(() => dispatch('increment')));
     assertLog(['Async action started [1]', 'Pending 0']);
     expect(container.textContent).toBe('Pending 0');
 
     // Dispatch a few more actions. None of these will start until the previous
     // one finishes.
-    await act(() => dispatch('increment'));
-    await act(() => dispatch('decrement'));
-    await act(() => dispatch('increment'));
+    await act(() => startTransition(() => dispatch('increment')));
+    await act(() => startTransition(() => dispatch('decrement')));
+    await act(() => startTransition(() => dispatch('increment')));
     assertLog([]);
 
     // Each action starts as soon as the previous one finishes.
@@ -1067,7 +1067,7 @@ describe('ReactDOMForm', () => {
 
     // Perform an action. This will increase the state by 1, as defined by the
     // stepSize prop.
-    await act(() => increment());
+    await act(() => startTransition(() => increment()));
     assertLog(['Pending 0', '1']);
 
     // Now increase the stepSize prop to 10. Subsequent steps will increase
@@ -1076,7 +1076,7 @@ describe('ReactDOMForm', () => {
     assertLog(['1']);
 
     // Increment again. The state should increase by 10.
-    await act(() => increment());
+    await act(() => startTransition(() => increment()));
     assertLog(['Pending 1', '11']);
   });
 
@@ -1113,11 +1113,11 @@ describe('ReactDOMForm', () => {
     await act(() => root.render(<App />));
     assertLog(['A']);
 
-    await act(() => action('B'));
+    await act(() => startTransition(() => action('B')));
     // The first dispatch will update the pending state.
     assertLog(['Pending A']);
-    await act(() => action('C'));
-    await act(() => action('D'));
+    await act(() => startTransition(() => action('C')));
+    await act(() => startTransition(() => action('D')));
     assertLog([]);
 
     await act(() => resolveText('B'));
@@ -1151,10 +1151,10 @@ describe('ReactDOMForm', () => {
 
       // Dispatch two actions. The first one is async, so it forces the second
       // one into an async queue.
-      await act(() => action('First action'));
+      await act(() => startTransition(() => action('First action')));
       assertLog(['Initial (pending)']);
       // This action won't run until the first one finishes.
-      await act(() => action('Second action'));
+      await act(() => startTransition(() => action('Second action')));
 
       // While the first action is still pending, update a prop. This causes the
       // inline action implementation to change, but it should not affect the
@@ -1169,7 +1169,9 @@ describe('ReactDOMForm', () => {
 
       // Confirm that if we dispatch yet another action, it uses the updated
       // action implementation.
-      await expect(act(() => action('Third action'))).rejects.toThrow('Oops!');
+      await expect(
+        act(() => startTransition(() => action('Third action'))),
+      ).rejects.toThrow('Oops!');
     },
   );
 
@@ -1192,7 +1194,7 @@ describe('ReactDOMForm', () => {
 
     // Perform an action. This will increase the state by 1, as defined by the
     // stepSize prop.
-    await act(() => increment());
+    await act(() => startTransition(() => increment()));
     assertLog(['Pending 0', '1']);
 
     // Now increase the stepSize prop to 10. Subsequent steps will increase
@@ -1201,7 +1203,7 @@ describe('ReactDOMForm', () => {
     assertLog(['1']);
 
     // Increment again. The state should increase by 10.
-    await act(() => increment());
+    await act(() => startTransition(() => increment()));
     assertLog(['Pending 1', '11']);
   });
 
@@ -1219,12 +1221,12 @@ describe('ReactDOMForm', () => {
     await act(() => root.render(<App />));
     assertLog(['A']);
 
-    await act(() => action(getText('B')));
+    await act(() => startTransition(() => action(getText('B'))));
     // The first dispatch will update the pending state.
     assertLog(['Pending A']);
-    await act(() => action('C'));
-    await act(() => action(getText('D')));
-    await act(() => action('E'));
+    await act(() => startTransition(() => action('C')));
+    await act(() => startTransition(() => action(getText('D'))));
+    await act(() => startTransition(() => action('E')));
     assertLog([]);
 
     await act(() => resolveText('B'));
@@ -1273,7 +1275,7 @@ describe('ReactDOMForm', () => {
     );
     assertLog(['A']);
 
-    await act(() => action('Oops!'));
+    await act(() => startTransition(() => action('Oops!')));
     assertLog([
       // Action begins, error has not thrown yet.
       'Pending A',
@@ -1290,8 +1292,8 @@ describe('ReactDOMForm', () => {
     // Trigger an error again, but this time, perform another action that
     // overrides the first one and fixes the error
     await act(() => {
-      action('Oops!');
-      action('B');
+      startTransition(() => action('Oops!'));
+      startTransition(() => action('B'));
     });
     assertLog(['Pending A', 'B']);
     expect(container.textContent).toBe('B');
@@ -1338,7 +1340,7 @@ describe('ReactDOMForm', () => {
     );
     assertLog(['A']);
 
-    await act(() => action('Oops!'));
+    await act(() => startTransition(() => action('Oops!')));
     // The first dispatch will update the pending state.
     assertLog(['Pending A']);
     await act(() => resolveText('Oops!'));
@@ -1352,8 +1354,8 @@ describe('ReactDOMForm', () => {
     // Trigger an error again, but this time, perform another action that
     // overrides the first one and fixes the error
     await act(() => {
-      action('Oops!');
-      action('B');
+      startTransition(() => action('Oops!'));
+      startTransition(() => action('B'));
     });
     assertLog(['Pending A']);
     await act(() => resolveText('B'));
@@ -1399,7 +1401,7 @@ describe('ReactDOMForm', () => {
     assertLog(['0']);
     expect(container.textContent).toBe('0');
 
-    await act(() => dispatch('increment'));
+    await act(() => startTransition(() => dispatch('increment')));
     assertLog(['Async action started [1]', 'Pending 0']);
     expect(container.textContent).toBe('Pending 0');
 
@@ -1408,6 +1410,77 @@ describe('ReactDOMForm', () => {
     expect(container.textContent).toBe('1');
   });
 
+  test('useActionState does not wrap action in a transition unless dispatch is in a transition', async () => {
+    let dispatch;
+    function App() {
+      const [state, _dispatch] = useActionState(() => {
+        return state + 1;
+      }, 0);
+      dispatch = _dispatch;
+      return <AsyncText text={'Count: ' + state} />;
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(() =>
+      root.render(
+        <Suspense fallback={<Text text="Loading..." />}>
+          <App />
+        </Suspense>,
+      ),
+    );
+    assertLog(['Suspend! [Count: 0]', 'Loading...']);
+    await act(() => resolveText('Count: 0'));
+    assertLog(['Count: 0']);
+
+    // Dispatch outside of a transition. This will trigger a loading state.
+    await act(() => dispatch());
+    assertLog(['Suspend! [Count: 1]', 'Loading...']);
+    expect(container.textContent).toBe('Loading...');
+
+    await act(() => resolveText('Count: 1'));
+    assertLog(['Count: 1']);
+    expect(container.textContent).toBe('Count: 1');
+
+    // Now dispatch inside of a transition. This one does not trigger a
+    // loading state.
+    await act(() => startTransition(() => dispatch()));
+    assertLog(['Count: 1', 'Suspend! [Count: 2]', 'Loading...']);
+    expect(container.textContent).toBe('Count: 1');
+
+    await act(() => resolveText('Count: 2'));
+    assertLog(['Count: 2']);
+    expect(container.textContent).toBe('Count: 2');
+  });
+
+  test('useActionState warns if async action is dispatched outside of a transition', async () => {
+    let dispatch;
+    function App() {
+      const [state, _dispatch] = useActionState(async () => {
+        return state + 1;
+      }, 0);
+      dispatch = _dispatch;
+      return <AsyncText text={'Count: ' + state} />;
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(() => root.render(<App />));
+    assertLog(['Suspend! [Count: 0]']);
+    await act(() => resolveText('Count: 0'));
+    assertLog(['Count: 0']);
+
+    // Dispatch outside of a transition.
+    await act(() => dispatch());
+    assertConsoleErrorDev([
+      [
+        'An async function was passed to useActionState, but it was ' +
+          'dispatched outside of an action context',
+        {withoutStack: true},
+      ],
+    ]);
+    assertLog(['Suspend! [Count: 1]']);
+    expect(container.textContent).toBe('Count: 0');
+  });
+
   test('uncontrolled form inputs are reset after the action completes', async () => {
     const formRef = React.createRef();
     const inputRef = React.createRef();

filename: packages/react-reconciler/src/ReactFiberHooks.js
patch: @@ -2006,65 +2006,87 @@ type ActionStateQueueNode<S, P> = {
   action: (Awaited<S>, P) => S,
   // This is never null because it's part of a circular linked list.
   next: ActionStateQueueNode<S, P>,
+
+  // Whether or not the action was dispatched as part of a transition. We use
+  // this to restore the transition context when the queued action is run. Once
+  // we're able to track parallel async actions, this should be updated to
+  // represent the specific transition instance the action is associated with.
+  isTransition: boolean,
+
+  // Implements the Thenable interface. We use it to suspend until the action
+  // finishes.
+  then: (listener: () => void) => void,
+  status: 'pending' | 'rejected' | 'fulfilled',
+  value: any,
+  reason: any,
+  listeners: Array<() => void>,
 };
 
 function dispatchActionState<S, P>(
   fiber: Fiber,
   actionQueue: ActionStateQueue<S, P>,
   setPendingState: boolean => void,
-  setState: Dispatch<S | Awaited<S>>,
+  setState: Dispatch<ActionStateQueueNode<S, P>>,
   payload: P,
 ): void {
   if (isRenderPhaseUpdate(fiber)) {
     throw new Error('Cannot update form state while rendering.');
   }
+
+  const actionNode: ActionStateQueueNode<S, P> = {
+    payload,
+    action: actionQueue.action,
+    next: (null: any), // circular
+
+    isTransition: true,
+
+    status: 'pending',
+    value: null,
+    reason: null,
+    listeners: [],
+    then(listener) {
+      // We know the only thing that subscribes to these promises is `use` so
+      // this implementation is simpler than a generic thenable. E.g. we don't
+      // bother to check if the thenable is still pending because `use` already
+      // does that.
+      actionNode.listeners.push(listener);
+    },
+  };
+
+  // Check if we're inside a transition. If so, we'll need to restore the
+  // transition context when the action is run.
+  const prevTransition = ReactSharedInternals.T;
+  if (prevTransition !== null) {
+    // Optimistically update the pending state, similar to useTransition.
+    // This will be reverted automatically when all actions are finished.
+    setPendingState(true);
+    // `actionNode` is a thenable that resolves to the return value of
+    // the action.
+    setState(actionNode);
+  } else {
+    // This is not a transition.
+    actionNode.isTransition = false;
+    setState(actionNode);
+  }
+
   const last = actionQueue.pending;
   if (last === null) {
     // There are no pending actions; this is the first one. We can run
     // it immediately.
-    const newLast: ActionStateQueueNode<S, P> = {
-      payload,
-      action: actionQueue.action,
-      next: (null: any), // circular
-    };
-    newLast.next = actionQueue.pending = newLast;
-
-    runActionStateAction(
-      actionQueue,
-      (setPendingState: any),
-      (setState: any),
-      newLast,
-    );
+    actionNode.next = actionQueue.pending = actionNode;
+    runActionStateAction(actionQueue, actionNode);
   } else {
     // There's already an action running. Add to the queue.
     const first = last.next;
-    const newLast: ActionStateQueueNode<S, P> = {
-      payload,
-      action: actionQueue.action,
-      next: first,
-    };
-    actionQueue.pending = last.next = newLast;
+    actionNode.next = first;
+    actionQueue.pending = last.next = actionNode;
   }
 }
 
 function runActionStateAction<S, P>(
   actionQueue: ActionStateQueue<S, P>,
-  setPendingState: boolean => void,
-  setState: Dispatch<S | Awaited<S>>,
   node: ActionStateQueueNode<S, P>,
 ) {
-  // This is a fork of startTransition
-  const prevTransition = ReactSharedInternals.T;
-  const currentTransition: BatchConfigTransition = {};
-  ReactSharedInternals.T = currentTransition;
-  if (__DEV__) {
-    ReactSharedInternals.T._updatedFibers = new Set();
-  }
-
-  // Optimistically update the pending state, similar to useTransition.
-  // This will be reverted automatically when all actions are finished.
-  setPendingState(true);
-
   // `node.action` represents the action function at the time it was dispatched.
   // If this action was queued, it might be stale, i.e. it's not necessarily the
   // most current implementation of the action, stored on `actionQueue`. This is
@@ -2074,93 +2096,106 @@ function runActionStateAction<S, P>(
   const action = node.action;
   const payload = node.payload;
   const prevState = actionQueue.state;
-  try {
-    const returnValue = action(prevState, payload);
-    const onStartTransitionFinish = ReactSharedInternals.S;
-    if (onStartTransitionFinish !== null) {
-      onStartTransitionFinish(currentTransition, returnValue);
+
+  if (node.isTransition) {
+    // The original dispatch was part of a transition. We restore its
+    // transition context here.
+
+    // This is a fork of startTransition
+    const prevTransition = ReactSharedInternals.T;
+    const currentTransition: BatchConfigTransition = {};
+    ReactSharedInternals.T = currentTransition;
+    if (__DEV__) {
+      ReactSharedInternals.T._updatedFibers = new Set();
     }
-    if (
-      returnValue !== null &&
-      typeof returnValue === 'object' &&
-      // $FlowFixMe[method-unbinding]
-      typeof returnValue.then === 'function'
-    ) {
-      const thenable = ((returnValue: any): Thenable<Awaited<S>>);
-
-      // Attach a listener to read the return state of the action. As soon as
-      // this resolves, we can run the next action in the sequence.
-      thenable.then(
-        (nextState: Awaited<S>) => {
-          actionQueue.state = nextState;
-          finishRunningActionStateAction(
-            actionQueue,
-            (setPendingState: any),
-            (setState: any),
-          );
-        },
-        () =>
-          finishRunningActionStateAction(
-            actionQueue,
-            (setPendingState: any),
-            (setState: any),
-          ),
-      );
+    try {
+      const returnValue = action(prevState, payload);
+      const onStartTransitionFinish = ReactSharedInternals.S;
+      if (onStartTransitionFinish !== null) {
+        onStartTransitionFinish(currentTransition, returnValue);
+      }
+      handleActionReturnValue(actionQueue, node, returnValue);
+    } catch (error) {
+      onActionError(actionQueue, node, error);
+    } finally {
+      ReactSharedInternals.T = prevTransition;
 
-      setState((thenable: any));
-    } else {
-      setState((returnValue: any));
-
-      const nextState = ((returnValue: any): Awaited<S>);
-      actionQueue.state = nextState;
-      finishRunningActionStateAction(
-        actionQueue,
-        (setPendingState: any),
-        (setState: any),
-      );
+      if (__DEV__) {
+        if (prevTransition === null && currentTransition._updatedFibers) {
+          const updatedFibersCount = currentTransition._updatedFibers.size;
+          currentTransition._updatedFibers.clear();
+          if (updatedFibersCount > 10) {
+            console.warn(
+              'Detected a large number of updates inside startTransition. ' +
+                'If this is due to a subscription please re-write it to use React provided hooks. ' +
+                'Otherwise concurrent mode guarantees are off the table.',
+            );
+          }
+        }
+      }
     }
-  } catch (error) {
-    // This is a trick to get the `useActionState` hook to rethrow the error.
-    // When it unwraps the thenable with the `use` algorithm, the error
-    // will be thrown.
-    const rejectedThenable: S = ({
-      then() {},
-      status: 'rejected',
-      reason: error,
-      // $FlowFixMe: Not sure why this doesn't work
-    }: RejectedThenable<Awaited<S>>);
-    setState(rejectedThenable);
-    finishRunningActionStateAction(
-      actionQueue,
-      (setPendingState: any),
-      (setState: any),
+  } else {
+    // The original dispatch was not part of a transition.
+    try {
+      const returnValue = action(prevState, payload);
+      handleActionReturnValue(actionQueue, node, returnValue);
+    } catch (error) {
+      onActionError(actionQueue, node, error);
+    }
+  }
+}
+
+function handleActionReturnValue<S, P>(
+  actionQueue: ActionStateQueue<S, P>,
+  node: ActionStateQueueNode<S, P>,
+  returnValue: mixed,
+) {
+  if (
+    returnValue !== null &&
+    typeof returnValue === 'object' &&
+    // $FlowFixMe[method-unbinding]
+    typeof returnValue.then === 'function'
+  ) {
+    const thenable = ((returnValue: any): Thenable<Awaited<S>>);
+    // Attach a listener to read the return state of the action. As soon as
+    // this resolves, we can run the next action in the sequence.
+    thenable.then(
+      (nextState: Awaited<S>) => {
+        onActionSuccess(actionQueue, node, nextState);
+      },
+      (error: mixed) => onActionError(actionQueue, node, error),
     );
-  } finally {
-    ReactSharedInternals.T = prevTransition;
 
     if (__DEV__) {
-      if (prevTransition === null && currentTransition._updatedFibers) {
-        const updatedFibersCount = currentTransition._updatedFibers.size;
-        currentTransition._updatedFibers.clear();
-        if (updatedFibersCount > 10) {
-          console.warn(
-            'Detected a large number of updates inside startTransition. ' +
-              'If this is due to a subscription please re-write it to use React provided hooks. ' +
-              'Otherwise concurrent mode guarantees are off the table.',
-          );
-        }
+      if (!node.isTransition) {
+        console.error(
+          'An async function was passed to useActionState, but it was ' +
+            'dispatched outside of an action context. This is likely not ' +
+            'what you intended. Either pass the dispatch function to an ' +
+            '`action` prop, or dispatch manually inside `startTransition`',
+        );
       }
     }
+  } else {
+    const nextState = ((returnValue: any): Awaited<S>);
+    onActionSuccess(actionQueue, node, nextState);
   }
 }
 
-function finishRunningActionStateAction<S, P>(
+function onActionSuccess<S, P>(
   actionQueue: ActionStateQueue<S, P>,
-  setPendingState: Dispatch<S | Awaited<S>>,
-  setState: Dispatch<S | Awaited<S>>,
+  actionNode: ActionStateQueueNode<S, P>,
+  nextState: Awaited<S>,
 ) {
-  // The action finished running. Pop it from the queue and run the next pending
-  // action, if there are any.
+  // The action finished running.
+  actionNode.status = 'fulfilled';
+  actionNode.value = nextState;
+  notifyActionListeners(actionNode);
+
+  actionQueue.state = nextState;
+
+  // Pop the action from the queue and run the next pending action, if there
+  // are any.
   const last = actionQueue.pending;
   if (last !== null) {
     const first = last.next;
@@ -2173,16 +2208,51 @@ function finishRunningActionStateAction<S, P>(
       last.next = next;
 
       // Run the next action.
-      runActionStateAction(
-        actionQueue,
-        (setPendingState: any),
-        (setState: any),
-        next,
-      );
+      runActionStateAction(actionQueue, next);
+    }
+  }
+}
+
+function onActionError<S, P>(
+  actionQueue: ActionStateQueue<S, P>,
+  actionNode: ActionStateQueueNode<S, P>,
+  error: mixed,
+) {
+  actionNode.status = 'rejected';
+  actionNode.reason = error;
+  notifyActionListeners(actionNode);
+
+  // Pop the action from the queue and run the next pending action, if there
+  // are any.
+  // TODO: We should instead abort all the remaining actions in the queue.
+  const last = actionQueue.pending;
+  if (last !== null) {
+    const first = last.next;
+    if (first === last) {
+      // This was the last action in the queue.
+      actionQueue.pending = null;
+    } else {
+      // Remove the first node from the circular queue.
+      const next = first.next;
+      last.next = next;
+
+      // Run the next action.
+      runActionStateAction(actionQueue, next);
     }
   }
 }
 
+function notifyActionListeners<S, P>(actionNode: ActionStateQueueNode<S, P>) {
+  // Notify React that the action has finished.
+  const listeners = actionNode.listeners;
+  for (let i = 0; i < listeners.length; i++) {
+    // This is always a React internal listener, so we don't need to worry
+    // about it throwing.
+    const listener = listeners[i];
+    listener();
+  }
+}
+
 function actionStateReducer<S>(oldState: S, newState: S): S {
   return newState;
 }