Fix: useTransition after use gets stuck in pending state (#29670)

When a component suspends with `use`, we switch to the "re-render"
dispatcher during the subsequent render attempt, so that we can reuse
the work from the initial attempt. However, once we run out of hooks
from the previous attempt, we should switch back to the regular "update"
dispatcher.

This is conceptually the same fix as the one introduced in
#26232. That fix only accounted
for initial mount, but the useTransition regression test added in
f829733 illustrates that we need to
handle updates, too.

The issue affects more than just useTransition but because most of the
behavior between the "re-render" and "update" dispatchers is the same
it's hard to contrive other scenarios in a test, which is probably why
it took so long for someone to notice.

Closes #28923 and #29209

---------

Co-authored-by: eps1lon <sebastian.silbermann@vercel.com>

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1083,20 +1083,49 @@ function useThenable<T>(thenable: Thenable<T>): T {
     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.
+
+  // When something suspends with `use`, we replay the component with the
+  // "re-render" dispatcher instead of the "mount" or "update" dispatcher.
+  //
+  // But if there are additional hooks that occur after the `use` invocation
+  // that suspended, they wouldn't have been processed during the previous
+  // attempt. So after we invoke `use` again, we may need to switch from the
+  // "re-render" dispatcher back to the "mount" or "update" dispatcher. That's
+  // what the following logic accounts for.
+  //
+  // TODO: Theoretically this logic only needs to go into the rerender
+  // dispatcher. Could optimize, but probably not be worth it.
+
+  // This is the same logic as in updateWorkInProgressHook.
+  const workInProgressFiber = currentlyRenderingFiber;
+  const nextWorkInProgressHook =
+    workInProgressHook === null
+      ? // We're at the beginning of the list, so read from the first hook from
+        // the fiber.
+        workInProgressFiber.memoizedState
+      : workInProgressHook.next;
+
+  if (nextWorkInProgressHook !== null) {
+    // There are still hooks remaining from the previous attempt.
+  } else {
+    // There are no remaining hooks from the previous attempt. We're no longer
+    // in "re-render" mode. Switch to the normal mount or update dispatcher.
+    //
+    // This is the same as the logic in renderWithHooks, except we don't bother
+    // to track the hook types debug information in this case (sufficient to
+    // only do that when nothing suspends).
+    const currentFiber = workInProgressFiber.alternate;
     if (__DEV__) {
-      ReactSharedInternals.H = HooksDispatcherOnMountInDEV;
+      if (currentFiber !== null && currentFiber.memoizedState !== null) {
+        ReactSharedInternals.H = HooksDispatcherOnUpdateInDEV;
+      } else {
+        ReactSharedInternals.H = HooksDispatcherOnMountInDEV;
+      }
     } else {
-      ReactSharedInternals.H = HooksDispatcherOnMount;
+      ReactSharedInternals.H =
+        currentFiber === null || currentFiber.memoizedState === null
+          ? HooksDispatcherOnMount
+          : HooksDispatcherOnUpdate;
     }
   }
   return result;

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

@@ -16,6 +16,7 @@ let act;
 let use;
 let useDebugValue;
 let useState;
+let useTransition;
 let useMemo;
 let useEffect;
 let Suspense;
@@ -38,6 +39,7 @@ describe('ReactUse', () => {
     use = React.use;
     useDebugValue = React.useDebugValue;
     useState = React.useState;
+    useTransition = React.useTransition;
     useMemo = React.useMemo;
     useEffect = React.useEffect;
     Suspense = React.Suspense;
@@ -1915,4 +1917,80 @@ describe('ReactUse', () => {
     assertLog(['Hi', 'World']);
     expect(root).toMatchRenderedOutput(<div>Hi World</div>);
   });
+
+  it(
+    'regression: does not get stuck in pending state after `use` suspends ' +
+      '(when `use` comes before all hooks)',
+    async () => {
+      // This is a regression test. The root cause was an issue where we failed to
+      // switch from the "re-render" dispatcher back to the "update" dispatcher
+      // after a `use` suspends and triggers a replay.
+      let update;
+      function App({promise}) {
+        const value = use(promise);
+
+        const [isPending, startLocalTransition] = useTransition();
+        update = () => {
+          startLocalTransition(() => {
+            root.render(<App promise={getAsyncText('Updated')} />);
+          });
+        };
+
+        return <Text text={value + (isPending ? ' (pending...)' : '')} />;
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(() => {
+        root.render(<App promise={Promise.resolve('Initial')} />);
+      });
+      assertLog(['Initial']);
+      expect(root).toMatchRenderedOutput('Initial');
+
+      await act(() => update());
+      assertLog(['Async text requested [Updated]', 'Initial (pending...)']);
+
+      await act(() => resolveTextRequests('Updated'));
+      assertLog(['Updated']);
+      expect(root).toMatchRenderedOutput('Updated');
+    },
+  );
+
+  it(
+    'regression: does not get stuck in pending state after `use` suspends ' +
+      '(when `use` in in the middle of hook list)',
+    async () => {
+      // Same as previous test but `use` comes in between two hooks.
+      let update;
+      function App({promise}) {
+        // This hook is only here to test that `use` resumes correctly after
+        // suspended even if it comes in between other hooks.
+        useState(false);
+
+        const value = use(promise);
+
+        const [isPending, startLocalTransition] = useTransition();
+        update = () => {
+          startLocalTransition(() => {
+            root.render(<App promise={getAsyncText('Updated')} />);
+          });
+        };
+
+        return <Text text={value + (isPending ? ' (pending...)' : '')} />;
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(() => {
+        root.render(<App promise={Promise.resolve('Initial')} />);
+      });
+      assertLog(['Initial']);
+      expect(root).toMatchRenderedOutput('Initial');
+
+      await act(() => update());
+      assertLog(['Async text requested [Updated]', 'Initial (pending...)']);
+
+      await act(() => resolveTextRequests('Updated'));
+      assertLog(['Updated']);
+      expect(root).toMatchRenderedOutput('Updated');
+    },
+  );
 });