Move renderDidSuspend logic to throwException

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -126,7 +126,6 @@ import {
   setShallowSuspenseListContext,
   ForceSuspenseFallback,
   setDefaultShallowSuspenseListContext,
-  isBadSuspenseFallback,
 } from './ReactFiberSuspenseContext';
 import {popHiddenContext} from './ReactFiberHiddenContext';
 import {findFirstSuspended} from './ReactFiberSuspenseComponent';
@@ -148,8 +147,6 @@ import {
   upgradeHydrationErrorsToRecoverable,
 } from './ReactFiberHydrationContext';
 import {
-  renderDidSuspend,
-  renderDidSuspendDelayIfPossible,
   renderHasNotSuspendedYet,
   getRenderTargetTime,
   getWorkInProgressTransitions,
@@ -1257,24 +1254,6 @@ function completeWork(
         if (nextDidTimeout) {
           const offscreenFiber: Fiber = (workInProgress.child: any);
           offscreenFiber.flags |= Visibility;
-
-          // TODO: This will still suspend a synchronous tree if anything
-          // in the concurrent tree already suspended during this render.
-          // This is a known bug.
-          if ((workInProgress.mode & ConcurrentMode) !== NoMode) {
-            // TODO: Move this back to throwException because this is too late
-            // if this is a large tree which is common for initial loads. We
-            // don't know if we should restart a render or not until we get
-            // this marker, and this is too late.
-            // If this render already had a ping or lower pri updates,
-            // and this is the first time we know we're going to suspend we
-            // should be able to immediately restart from within throwException.
-            if (isBadSuspenseFallback(current, newProps)) {
-              renderDidSuspendDelayIfPossible();
-            } else {
-              renderDidSuspend();
-            }
-          }
         }
       }
 

packages/react-reconciler/src/ReactFiberSuspenseContext.js

@@ -62,6 +62,11 @@ export function isBadSuspenseFallback(
   // Check if this is a "bad" fallback state or a good one. A bad fallback state
   // is one that we only show as a last resort; if this is a transition, we'll
   // block it from displaying, and wait for more data to arrive.
+  // TODO: This is function is only used by the `use` implementation. There's
+  // identical logic in `throwException`, and also in the begin phase of the
+  // Suspense component. Since we're already computing this in the begin phase,
+  // track it on stack instead of re-computing it on demand. Less code, less
+  // duplicated logic, less computation.
   if (current !== null) {
     const prevState: SuspenseState = current.memoizedState;
     const isShowingFallback = prevState !== null;

packages/react-reconciler/src/ReactFiberThrow.js

@@ -13,6 +13,7 @@ import type {CapturedValue} from './ReactCapturedValue';
 import type {Update} from './ReactFiberClassUpdateQueue';
 import type {Wakeable} from 'shared/ReactTypes';
 import type {OffscreenQueue} from './ReactFiberOffscreenComponent';
+import type {SuspenseState} from './ReactFiberSuspenseComponent';
 
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
 import {
@@ -39,6 +40,7 @@ import {
   enableDebugTracing,
   enableLazyContextPropagation,
   enableUpdaterTracking,
+  enableSuspenseAvoidThisFallback,
 } from 'shared/ReactFeatureFlags';
 import {createCapturedValueAtFiber} from './ReactCapturedValue';
 import {
@@ -58,6 +60,7 @@ import {
   isAlreadyFailedLegacyErrorBoundary,
   attachPingListener,
   restorePendingUpdaters,
+  renderDidSuspend,
 } from './ReactFiberWorkLoop';
 import {propagateParentContextChangesToDeferredTree} from './ReactFiberNewContext';
 import {logCapturedError} from './ReactFiberErrorLogger';
@@ -349,11 +352,60 @@ function throwException(
       }
     }
 
-    // Schedule the nearest Suspense to re-render the timed out view.
+    // Mark the nearest Suspense boundary to switch to rendering a fallback.
     const suspenseBoundary = getSuspenseHandler();
     if (suspenseBoundary !== null) {
       switch (suspenseBoundary.tag) {
         case SuspenseComponent: {
+          // If this suspense boundary is not already showing a fallback, mark
+          // the in-progress render as suspended. We try to perform this logic
+          // as soon as soon as possible during the render phase, so the work
+          // loop can know things like whether it's OK to switch to other tasks,
+          // or whether it can wait for data to resolve before continuing.
+          // TODO: Most of these checks are already performed when entering a
+          // Suspense boundary. We should track the information on the stack so
+          // we don't have to recompute it on demand. This would also allow us
+          // to unify with `use` which needs to perform this logic even sooner,
+          // before `throwException` is called.
+          if (sourceFiber.mode & ConcurrentMode) {
+            if (getIsHydrating()) {
+              // A dehydrated boundary is considered a fallback state. We don't
+              // have to suspend.
+            } else {
+              const current = suspenseBoundary.alternate;
+              if (current === null) {
+                // This is a new mount. Unless this is an "avoided" fallback
+                // (experimental feature) this should not delay the tree
+                // from appearing.
+                const nextProps = suspenseBoundary.pendingProps;
+                if (
+                  enableSuspenseAvoidThisFallback &&
+                  nextProps.unstable_avoidThisFallback === true
+                ) {
+                  // Experimental feature: Some fallbacks are always bad
+                  renderDidSuspendDelayIfPossible();
+                } else {
+                  // Show a fallback without delaying. The only reason we mark
+                  // this case at all is so we can throttle the appearance of
+                  // new fallbacks. If we did nothing here, all other behavior
+                  // would be the same, except it wouldn't throttle.
+                  renderDidSuspend();
+                }
+              } else {
+                const prevState: SuspenseState = current.memoizedState;
+                if (prevState !== null) {
+                  // This boundary is currently showing a fallback. Don't need
+                  // to suspend.
+                } else {
+                  // This boundary is currently showing content. Switching to a
+                  // fallback will cause that content to disappear. Tell the
+                  // work loop to delay the commit, if possible.
+                  renderDidSuspendDelayIfPossible();
+                }
+              }
+            }
+          }
+
           suspenseBoundary.flags &= ~ForceClientRender;
           markSuspenseBoundaryShouldCapture(
             suspenseBoundary,

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -1851,9 +1851,10 @@ function handleThrow(root, thrownValue): void {
 }
 
 function shouldAttemptToSuspendUntilDataResolves() {
-  // TODO: We should be able to move the
-  // renderDidSuspend/renderDidSuspendDelayIfPossible logic into this function,
-  // instead of repeating it in the complete phase. Or something to that effect.
+  // TODO: This function needs to have parity with
+  // renderDidSuspendDelayIfPossible, but it currently doesn't. (It only affects
+  // the `use` API.) Fix by unifying the logic here with the equivalent checks
+  // in `throwException` and in the begin phase of Suspense.
 
   if (includesOnlyRetries(workInProgressRootRenderLanes)) {
     // We can always wait during a retry.
@@ -3366,8 +3367,16 @@ function pingSuspendedRoot(
         includesOnlyRetries(workInProgressRootRenderLanes) &&
         now() - globalMostRecentFallbackTime < FALLBACK_THROTTLE_MS)
     ) {
-      // Restart from the root.
-      prepareFreshStack(root, NoLanes);
+      // Force a restart from the root by unwinding the stack. Unless this is
+      // being called from the render phase, because that would cause a crash.
+      if ((executionContext & RenderContext) === NoContext) {
+        prepareFreshStack(root, NoLanes);
+      } else {
+        // TODO: If this does happen during the render phase, we should throw
+        // the special internal exception that we use to interrupt the stack for
+        // selective hydration. That was temporarily reverted but we once we add
+        // it back we can use it here.
+      }
     } else {
       // Even though we can't restart right now, we might get an
       // opportunity later. So we mark this render as having a ping.

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

@@ -995,14 +995,10 @@ describe('ReactSuspenseWithNoopRenderer', () => {
     expect(Scheduler).toFlushAndYieldThrough(['Suspend! [Async]', 'Sibling']);
 
     await resolveText('Async');
-    expect(Scheduler).toFlushAndYield([
-      // We've now pinged the boundary but we don't know if we should restart yet,
-      // because we haven't completed the suspense boundary.
-      'Loading...',
-      // Once we've completed the boundary we restarted.
-      'Async',
-      'Sibling',
-    ]);
+
+    // Because we're already showing a fallback, interrupt the current render
+    // and restart immediately.
+    expect(Scheduler).toFlushAndYield(['Async', 'Sibling']);
     expect(root).toMatchRenderedOutput(
       <>
         <span prop="Async" />
@@ -4218,4 +4214,80 @@ describe('ReactSuspenseWithNoopRenderer', () => {
     });
     expect(Scheduler).toHaveYielded(['Unmount Child']);
   });
+
+  // @gate enableLegacyCache
+  it(
+    'regression test: pinging synchronously within the render phase ' +
+      'does not unwind the stack',
+    async () => {
+      // This is a regression test that reproduces a very specific scenario that
+      // used to cause a crash.
+      const thenable = {
+        then(resolve) {
+          resolve('hi');
+        },
+        status: 'pending',
+      };
+
+      function ImmediatelyPings() {
+        if (thenable.status === 'pending') {
+          thenable.status = 'fulfilled';
+          throw thenable;
+        }
+        return <Text text="Hi" />;
+      }
+
+      function App({showMore}) {
+        return (
+          <div>
+            <Suspense fallback={<Text text="Loading..." />}>
+              {showMore ? (
+                <>
+                  <AsyncText text="Async" />
+                </>
+              ) : null}
+            </Suspense>
+            {showMore ? (
+              <Suspense>
+                <ImmediatelyPings />
+              </Suspense>
+            ) : null}
+          </div>
+        );
+      }
+
+      // Initial render. This mounts a Suspense boundary, so that in the next
+      // update we can trigger a "suspend with delay" scenario.
+      const root = ReactNoop.createRoot();
+      await act(async () => {
+        root.render(<App showMore={false} />);
+      });
+      expect(Scheduler).toHaveYielded([]);
+      expect(root).toMatchRenderedOutput(<div />);
+
+      // Update. This will cause two separate trees to suspend. The first tree
+      // will be inside an already mounted Suspense boundary, so it will trigger
+      // a "suspend with delay". The second tree will be a new Suspense
+      // boundary, but the thenable that is thrown will immediately call its
+      // ping listener.
+      //
+      // Before the bug was fixed, this would lead to a `prepareFreshStack` call
+      // that unwinds the work-in-progress stack. When that code was written, it
+      // was expected that pings always happen from an asynchronous task (or
+      // microtask). But this test shows an example where that's not the case.
+      //
+      // The fix was to check if we're in the render phase before calling
+      // `prepareFreshStack`.
+      await act(async () => {
+        root.render(<App showMore={true} />);
+      });
+      expect(Scheduler).toHaveYielded(['Suspend! [Async]', 'Loading...', 'Hi']);
+      expect(root).toMatchRenderedOutput(
+        <div>
+          <span prop="Loading..." />
+          <span prop="Hi" />
+        </div>,
+      );
+    },
+  );
 });