From 436f912a021f1c24cf5628c3c96cc5c3f1878b2b Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Tue, 9 Apr 2024 17:11:46 -0400
Subject: [PATCH] Fix: Suspend while recovering from hydration error (#28800)

Fixes a bug that happens when an error occurs during hydration, React
switches to client rendering, and then the client render suspends. It
works correctly if there's a Suspense boundary on the stack, but not if
it happens in the shell of the app.

Prior to this fix, the app would crash with an "Unknown root exit
status" error.

I left a TODO comment for how we might refactor this code to be less
confusing in the future.
---
 .../ReactDOMFizzShellHydration-test.js        | 75 +++++++++++++++++++
 .../src/ReactFiberWorkLoop.js                 | 19 ++++-
 2 files changed, 91 insertions(+), 3 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzShellHydration-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzShellHydration-test.js
index f4df63ca7a652..999fbd9f7190d 100644
--- a/packages/react-dom/src/__tests__/ReactDOMFizzShellHydration-test.js
+++ b/packages/react-dom/src/__tests__/ReactDOMFizzShellHydration-test.js
@@ -548,4 +548,79 @@ describe('ReactDOMFizzShellHydration', () => {
     ]);
     expect(container.textContent).toBe('Hello world');
   });
+
+  it(
+    'handles suspending while recovering from a hydration error (in the ' +
+      'shell, no Suspense boundary)',
+    async () => {
+      const useSyncExternalStore = React.useSyncExternalStore;
+
+      let isClient = false;
+
+      let resolve;
+      const clientPromise = new Promise(res => {
+        resolve = res;
+      });
+
+      function App() {
+        const state = useSyncExternalStore(
+          function subscribe() {
+            return () => {};
+          },
+          function getSnapshot() {
+            return 'Client';
+          },
+          function getServerSnapshot() {
+            const isHydrating = isClient;
+            if (isHydrating) {
+              // This triggers an error during hydration
+              throw new Error('Oops!');
+            }
+            return 'Server';
+          },
+        );
+
+        if (state === 'Client') {
+          return React.use(clientPromise);
+        }
+
+        return state;
+      }
+
+      // Server render
+      await serverAct(async () => {
+        const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<App />);
+        pipe(writable);
+      });
+      assertLog([]);
+
+      expect(container.innerHTML).toBe('Server');
+
+      // During hydration, an error is thrown. React attempts to recover by
+      // switching to client render
+      isClient = true;
+      await clientAct(async () => {
+        ReactDOMClient.hydrateRoot(container, <App />, {
+          onRecoverableError(error) {
+            Scheduler.log('onRecoverableError: ' + error.message);
+            if (error.cause) {
+              Scheduler.log('Cause: ' + error.cause.message);
+            }
+          },
+        });
+      });
+      expect(container.innerHTML).toBe('Server'); // Still suspended
+      assertLog([]);
+
+      await clientAct(async () => {
+        resolve('Client');
+      });
+      assertLog([
+        'onRecoverableError: There was an error while hydrating but React was ' +
+          'able to recover by instead client rendering the entire root.',
+        'Cause: Oops!',
+      ]);
+      expect(container.innerHTML).toBe('Client');
+    },
+  );
 });
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.js b/packages/react-reconciler/src/ReactFiberWorkLoop.js
index 17c1760755a7f..8b0c8729506cc 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.js
+++ b/packages/react-reconciler/src/ReactFiberWorkLoop.js
@@ -931,19 +931,32 @@ export function performConcurrentWorkOnRoot(
 
         // Check if something threw
         if (exitStatus === RootErrored) {
-          const originallyAttemptedLanes = lanes;
+          const lanesThatJustErrored = lanes;
           const errorRetryLanes = getLanesToRetrySynchronouslyOnError(
             root,
-            originallyAttemptedLanes,
+            lanesThatJustErrored,
           );
           if (errorRetryLanes !== NoLanes) {
             lanes = errorRetryLanes;
             exitStatus = recoverFromConcurrentError(
               root,
-              originallyAttemptedLanes,
+              lanesThatJustErrored,
               errorRetryLanes,
             );
             renderWasConcurrent = false;
+            // Need to check the exit status again.
+            if (exitStatus !== RootErrored) {
+              // The root did not error this time. Restart the exit algorithm
+              // from the beginning.
+              // TODO: Refactor the exit algorithm to be less confusing. Maybe
+              // more branches + recursion instead of a loop. I think the only
+              // thing that causes it to be a loop is the RootDidNotComplete
+              // check. If that's true, then we don't need a loop/recursion
+              // at all.
+              continue;
+            } else {
+              // The root errored yet again. Proceed to commit the tree.
+            }
           }
         }
         if (exitStatus === RootFatalErrored) {