Bugfix: Fix crash when suspending in shell during useSES re-render (facebook#27199)

This adds a regression test for a bug where, after a store mutation, the
updated data causes the shell of the app to suspend.

When re-rendering due to a concurrent store mutation, we must check for
the RootDidNotComplete exit status again.

As a follow-up, I'm going to look into to cleaning up the places where
we check the exit status, so bugs like these are less likely. I think we
might be able to combine most of it into a single switch statement.

Author: acdlite

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -892,58 +892,39 @@ export function performConcurrentWorkOnRoot(
   let exitStatus = shouldTimeSlice
     ? renderRootConcurrent(root, lanes)
     : renderRootSync(root, lanes);
-  if (exitStatus !== RootInProgress) {
-    if (exitStatus === RootErrored) {
-      // If something threw an error, try rendering one more time. We'll
-      // render synchronously to block concurrent data mutations, and we'll
-      // includes all pending updates are included. If it still fails after
-      // the second attempt, we'll give up and commit the resulting tree.
-      const originallyAttemptedLanes = lanes;
-      const errorRetryLanes = getLanesToRetrySynchronouslyOnError(
-        root,
-        originallyAttemptedLanes,
-      );
-      if (errorRetryLanes !== NoLanes) {
-        lanes = errorRetryLanes;
-        exitStatus = recoverFromConcurrentError(
-          root,
-          originallyAttemptedLanes,
-          errorRetryLanes,
-        );
-      }
-    }
-    if (exitStatus === RootFatalErrored) {
-      const fatalError = workInProgressRootFatalError;
-      prepareFreshStack(root, NoLanes);
-      markRootSuspended(root, lanes);
-      ensureRootIsScheduled(root);
-      throw fatalError;
-    }
 
-    if (exitStatus === RootDidNotComplete) {
-      // The render unwound without completing the tree. This happens in special
-      // cases where need to exit the current render without producing a
-      // consistent tree or committing.
-      markRootSuspended(root, lanes);
-    } else {
-      // The render completed.
-
-      // Check if this render may have yielded to a concurrent event, and if so,
-      // confirm that any newly rendered stores are consistent.
-      // TODO: It's possible that even a concurrent render may never have yielded
-      // to the main thread, if it was fast enough, or if it expired. We could
-      // skip the consistency check in that case, too.
-      const renderWasConcurrent = !includesBlockingLane(root, lanes);
-      const finishedWork: Fiber = (root.current.alternate: any);
-      if (
-        renderWasConcurrent &&
-        !isRenderConsistentWithExternalStores(finishedWork)
-      ) {
-        // A store was mutated in an interleaved event. Render again,
-        // synchronously, to block further mutations.
-        exitStatus = renderRootSync(root, lanes);
+  if (exitStatus !== RootInProgress) {
+    let renderWasConcurrent = shouldTimeSlice;
+    do {
+      if (exitStatus === RootDidNotComplete) {
+        // The render unwound without completing the tree. This happens in special
+        // cases where need to exit the current render without producing a
+        // consistent tree or committing.
+        markRootSuspended(root, lanes);
+      } else {
+        // The render completed.
+
+        // Check if this render may have yielded to a concurrent event, and if so,
+        // confirm that any newly rendered stores are consistent.
+        // TODO: It's possible that even a concurrent render may never have yielded
+        // to the main thread, if it was fast enough, or if it expired. We could
+        // skip the consistency check in that case, too.
+        const finishedWork: Fiber = (root.current.alternate: any);
+        if (
+          renderWasConcurrent &&
+          !isRenderConsistentWithExternalStores(finishedWork)
+        ) {
+          // A store was mutated in an interleaved event. Render again,
+          // synchronously, to block further mutations.
+          exitStatus = renderRootSync(root, lanes);
+          // We assume the tree is now consistent because we didn't yield to any
+          // concurrent events.
+          renderWasConcurrent = false;
+          // Need to check the exit status again.
+          continue;
+        }
 
-        // We need to check again if something threw
+        // Check if something threw
         if (exitStatus === RootErrored) {
           const originallyAttemptedLanes = lanes;
           const errorRetryLanes = getLanesToRetrySynchronouslyOnError(
@@ -957,8 +938,7 @@ export function performConcurrentWorkOnRoot(
               originallyAttemptedLanes,
               errorRetryLanes,
             );
-            // We assume the tree is now consistent because we didn't yield to any
-            // concurrent events.
+            renderWasConcurrent = false;
           }
         }
         if (exitStatus === RootFatalErrored) {
@@ -969,16 +949,14 @@ export function performConcurrentWorkOnRoot(
           throw fatalError;
         }
 
-        // FIXME: Need to check for RootDidNotComplete again. The factoring here
-        // isn't ideal.
+        // We now have a consistent tree. The next step is either to commit it,
+        // or, if something suspended, wait to commit it after a timeout.
+        root.finishedWork = finishedWork;
+        root.finishedLanes = lanes;
+        finishConcurrentRender(root, exitStatus, finishedWork, lanes);
       }
-
-      // We now have a consistent tree. The next step is either to commit it,
-      // or, if something suspended, wait to commit it after a timeout.
-      root.finishedWork = finishedWork;
-      root.finishedLanes = lanes;
-      finishConcurrentRender(root, exitStatus, finishedWork, lanes);
-    }
+      break;
+    } while (true);
   }
 
   ensureRootIsScheduled(root);

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

@@ -19,6 +19,7 @@ let forwardRef;
 let useImperativeHandle;
 let useRef;
 let useState;
+let use;
 let startTransition;
 let waitFor;
 let waitForAll;
@@ -41,6 +42,7 @@ describe('useSyncExternalStore', () => {
     forwardRef = React.forwardRef;
     useRef = React.useRef;
     useState = React.useState;
+    use = React.use;
     useSyncExternalStore = React.useSyncExternalStore;
     startTransition = React.startTransition;
 
@@ -212,4 +214,84 @@ describe('useSyncExternalStore', () => {
     });
     assertLog(['value:initial']);
   });
+
+  test(
+    'regression: suspending in shell after synchronously patching ' +
+      'up store mutation',
+    async () => {
+      // Tests a case where a store is mutated during a concurrent event, then
+      // during the sync re-render, a synchronous render is triggered.
+
+      const store = createExternalStore('Initial');
+
+      let resolve;
+      const promise = new Promise(r => {
+        resolve = r;
+      });
+
+      function A() {
+        const value = useSyncExternalStore(store.subscribe, store.getState);
+
+        if (value === 'Updated') {
+          try {
+            use(promise);
+          } catch (x) {
+            Scheduler.log('Suspend A');
+            throw x;
+          }
+        }
+
+        return <Text text={'A: ' + value} />;
+      }
+
+      function B() {
+        const value = useSyncExternalStore(store.subscribe, store.getState);
+        return <Text text={'B: ' + value} />;
+      }
+
+      function App() {
+        return (
+          <>
+            <span>
+              <A />
+            </span>
+            <span>
+              <B />
+            </span>
+          </>
+        );
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(async () => {
+        // A and B both read from the same store. Partially render A.
+        startTransition(() => root.render(<App />));
+        // A reads the initial value of the store.
+        await waitFor(['A: Initial']);
+
+        // Before B renders, mutate the store.
+        store.set('Updated');
+      });
+      assertLog([
+        // B reads the updated value of the store.
+        'B: Updated',
+        // This should a synchronous re-render of A using the updated value. In
+        // this test, this causes A to suspend.
+        'Suspend A',
+      ]);
+      // Nothing has committed, because A suspended and no fallback
+      // was provided.
+      expect(root).toMatchRenderedOutput(null);
+
+      // Resolve the data and finish rendering.
+      await act(() => resolve());
+      assertLog(['A: Updated', 'B: Updated']);
+      expect(root).toMatchRenderedOutput(
+        <>
+          <span>A: Updated</span>
+          <span>B: Updated</span>
+        </>,
+      );
+    },
+  );
 });