Batch async updates to classes, roots too

This implements the previous fix for class components and root-level
updates, too. I put this in its own step because there are some extra
changes related to avoiding a context stack mismatch.

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -137,6 +137,7 @@ import {
   cloneUpdateQueue,
   initializeUpdateQueue,
   enqueueCapturedUpdate,
+  suspendIfUpdateReadFromEntangledAsyncAction,
 } from './ReactFiberClassUpdateQueue';
 import {
   NoLane,
@@ -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;
@@ -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;

packages/react-reconciler/src/ReactFiberClassComponent.js

@@ -53,6 +53,7 @@ import {
   ForceUpdate,
   initializeUpdateQueue,
   cloneUpdateQueue,
+  suspendIfUpdateReadFromEntangledAsyncAction,
 } from './ReactFiberClassUpdateQueue';
 import {NoLanes} from './ReactFiberLane';
 import {
@@ -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;
   }
 
@@ -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 &&
@@ -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 (

packages/react-reconciler/src/ReactFiberClassUpdateQueue.js

@@ -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,
@@ -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);
 
@@ -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

packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js

@@ -1461,4 +1461,184 @@ describe('ReactAsyncActions', () => {
       expect(root).toMatchRenderedOutput(<span>C</span>);
     },
   );
+
+  // @gate enableAsyncActions
+  test(
+    'updates in an async action are entangled even if useTransition hook ' +
+      'is unmounted before it finishes (class component)',
+    async () => {
+      let startTransition;
+      function Updater() {
+        const [isPending, _start] = useTransition();
+        startTransition = _start;
+        return (
+          <span>
+            <Text text={'Pending: ' + isPending} />
+          </span>
+        );
+      }
+
+      let setText;
+      class Sibling extends React.Component {
+        state = {text: 'A'};
+        render() {
+          setText = text => this.setState({text});
+          return (
+            <span>
+              <Text text={this.state.text} />
+            </span>
+          );
+        }
+      }
+
+      function App({showUpdater}) {
+        return (
+          <>
+            {showUpdater ? <Updater /> : null}
+            <Sibling />
+          </>
+        );
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(() => {
+        root.render(<App showUpdater={true} />);
+      });
+      assertLog(['Pending: false', 'A']);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span>Pending: false</span>
+          <span>A</span>
+        </>,
+      );
+
+      // Start an async action that has multiple updates with async
+      // operations in between.
+      await act(() => {
+        startTransition(async () => {
+          Scheduler.log('Async action started');
+          startTransition(() => setText('B'));
+
+          await getText('Wait before updating to C');
+
+          Scheduler.log('Async action ended');
+          startTransition(() => setText('C'));
+        });
+      });
+      assertLog(['Async action started', 'Pending: true']);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span>Pending: true</span>
+          <span>A</span>
+        </>,
+      );
+
+      // Delete the component that contains the useTransition hook. This
+      // component no longer blocks the transition from completing. But the
+      // pending update to Sibling should not be allowed to finish, because it's
+      // part of the async action.
+      await act(() => {
+        root.render(<App showUpdater={false} />);
+      });
+      assertLog(['A']);
+      expect(root).toMatchRenderedOutput(<span>A</span>);
+
+      // Finish the async action. Notice the intermediate B state was never
+      // shown, because it was batched with the update that came later in the
+      // same action.
+      await act(() => resolveText('Wait before updating to C'));
+      assertLog(['Async action ended', 'C']);
+      expect(root).toMatchRenderedOutput(<span>C</span>);
+
+      // Check that subsequent updates are unaffected.
+      await act(() => setText('D'));
+      assertLog(['D']);
+      expect(root).toMatchRenderedOutput(<span>D</span>);
+    },
+  );
+
+  // @gate enableAsyncActions
+  test(
+    'updates in an async action are entangled even if useTransition hook ' +
+      'is unmounted before it finishes (root update)',
+    async () => {
+      let startTransition;
+      function Updater() {
+        const [isPending, _start] = useTransition();
+        startTransition = _start;
+        return (
+          <span>
+            <Text text={'Pending: ' + isPending} />
+          </span>
+        );
+      }
+
+      let setShowUpdater;
+      function App({text}) {
+        const [showUpdater, _setShowUpdater] = useState(true);
+        setShowUpdater = _setShowUpdater;
+        return (
+          <>
+            {showUpdater ? <Updater /> : null}
+            <span>
+              <Text text={text} />
+            </span>
+          </>
+        );
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(() => {
+        root.render(<App text="A" />);
+      });
+      assertLog(['Pending: false', 'A']);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span>Pending: false</span>
+          <span>A</span>
+        </>,
+      );
+
+      // Start an async action that has multiple updates with async
+      // operations in between.
+      await act(() => {
+        startTransition(async () => {
+          Scheduler.log('Async action started');
+          startTransition(() => root.render(<App text="B" />));
+
+          await getText('Wait before updating to C');
+
+          Scheduler.log('Async action ended');
+          startTransition(() => root.render(<App text="C" />));
+        });
+      });
+      assertLog(['Async action started', 'Pending: true']);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span>Pending: true</span>
+          <span>A</span>
+        </>,
+      );
+
+      // Delete the component that contains the useTransition hook. This
+      // component no longer blocks the transition from completing. But the
+      // pending update to Sibling should not be allowed to finish, because it's
+      // part of the async action.
+      await act(() => setShowUpdater(false));
+      assertLog(['A']);
+      expect(root).toMatchRenderedOutput(<span>A</span>);
+
+      // Finish the async action. Notice the intermediate B state was never
+      // shown, because it was batched with the update that came later in the
+      // same action.
+      await act(() => resolveText('Wait before updating to C'));
+      assertLog(['Async action ended', 'C']);
+      expect(root).toMatchRenderedOutput(<span>C</span>);
+
+      // Check that subsequent updates are unaffected.
+      await act(() => root.render(<App text="D" />));
+      assertLog(['D']);
+      expect(root).toMatchRenderedOutput(<span>D</span>);
+    },
+  );
 });