Move renderDidSuspend logic to throwException

This is part of a larger refactor I'm doing to simplify how this logic
works, since it's currently spread out and duplicated across several
different parts of the work loop.

When a Suspense boundary shows a fallback, and it wasn't already showing
a fallback, we mark the render as suspended. Knowing whether the
in-progress tree includes a new fallback is useful for several reasons:
we can choose to interrupt the current render and switch to a different
task, we can suspend the work loop until more data becomes available, we
can throttle the appearance of multiple fallback states, and so on.

Currently this logic happens in the complete phase, but because we use
renderDidSuspendWithDelay as a signal to interrupt the work loop, it's
best to do it early as we possibly can.

A subsequent step will move the logic even earlier, as soon as we
attempt to unwrap an unresolved promise, so that `use` can determine
whether it's OK to pause the entire work loop and wait for the promise.
There's some existing code that attempts to do this but it doesn't have
parity with how `renderDidSuspend` works, which is the main motivation
for this series of refactors.

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.

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" />