Unify use and renderDidSuspendWithDelay impl

acdlite · acdlite · commit 06f075d003e1 · 2022-12-22T00:57:00.000-05:00
When unwrapping a promise with `use`, we sometimes suspend the work loop
from rendering anything else until the data has resolved. This is
different from how Suspense works in the old throw-a-promise world,
where rather than suspend rendering midway through the render phase, we
prepare a fallback and block the commit at the end, if necessary;
however, the logic for determining whether it's OK to block is the same.
The implementation is only incidentally different because it happens in
two different parts of the code. This means for `use`, we end up doing
the same checks twice, which is wasteful in terms of computataion, but
also introduces a risk that the logic will accidentally diverage.

This unifies the implementation by moving it into the SuspenseContext
module. Most of the logic for deciding whether to suspend is already
performed in the begin phase of SuspenseComponent, so it makes sense to
store that information on the stack rather than recompute it on demand.

The way I've chosen to model this is to track whether the work loop is
rendering inside the "shell" of the tree. The shell is defined as the
part of the tree that's visible in the current UI. Once we enter a 
new Suspense boundary (or a hidden Offscreen boundary, which acts a
Suspense boundary), we're no longer in the shell. This is already how
Suspense behavior was modeled in terms of UX, so using this concept 
directly in the implementation turns out to result in less code
than before.

For the most part, this is purely an internal refactor, though it does
fix a bug in the `use` implementation related to nested Suspense
boundaries. I wouldn't be surprised if it happens to fix other bugs that
we haven't yet discovered, especially around Offscreen. I'll add more
tests as I think of them.
diff --git a/packages/react-reconciler/src/ReactFiberSuspenseContext.js b/packages/react-reconciler/src/ReactFiberSuspenseContext.js
@@ -9,7 +9,8 @@
 
 import type {Fiber} from './ReactInternalTypes';
 import type {StackCursor} from './ReactFiberStack';
-import type {SuspenseState, SuspenseProps} from './ReactFiberSuspenseComponent';
+import type {SuspenseProps, SuspenseState} from './ReactFiberSuspenseComponent';
+import type {OffscreenState} from './ReactFiberOffscreenComponent';
 
 import {enableSuspenseAvoidThisFallback} from 'shared/ReactFeatureFlags';
 import {createCursor, push, pop} from './ReactFiberStack';
@@ -22,86 +23,72 @@ const suspenseHandlerStackCursor: StackCursor<Fiber | null> = createCursor(
   null,
 );
 
-function shouldAvoidedBoundaryCapture(
-  workInProgress: Fiber,
-  handlerOnStack: Fiber,
-  props: any,
-): boolean {
-  if (enableSuspenseAvoidThisFallback) {
-    // If the parent is already showing content, and we're not inside a hidden
-    // tree, then we should show the avoided fallback.
-    if (handlerOnStack.alternate !== null && !isCurrentTreeHidden()) {
-      return true;
-    }
-
-    // If the handler on the stack is also an avoided boundary, then we should
-    // favor this inner one.
-    if (
-      handlerOnStack.tag === SuspenseComponent &&
-      handlerOnStack.memoizedProps.unstable_avoidThisFallback === true
-    ) {
-      return true;
-    }
-
-    // If this avoided boundary is dehydrated, then it should capture.
-    const suspenseState: SuspenseState | null = workInProgress.memoizedState;
-    if (suspenseState !== null && suspenseState.dehydrated !== null) {
-      return true;
-    }
-  }
-
-  // If none of those cases apply, then we should avoid this fallback and show
-  // the outer one instead.
-  return false;
-}
-
-export function isBadSuspenseFallback(
-  current: Fiber | null,
-  nextProps: SuspenseProps,
-): boolean {
-  // 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;
-    if (!isShowingFallback && !isCurrentTreeHidden()) {
-      // It's bad to switch to a fallback if content is already visible
-      return true;
-    }
-  }
-
-  if (
-    enableSuspenseAvoidThisFallback &&
-    nextProps.unstable_avoidThisFallback === true
-  ) {
-    // Experimental: Some fallbacks are always bad
-    return true;
-  }
-
-  return false;
+// Represents the outermost boundary that is not visible in the current tree.
+// Everything above this is the "shell". When this is null, it means we're
+// rendering in the shell of the app. If it's non-null, it means we're rendering
+// deeper than the shell, inside a new tree that wasn't already visible.
+//
+// The main way we use this concept is to determine whether showing a fallback
+// would result in a desirable or undesirable loading state. Activing a fallback
+// in the shell is considered an undersirable loading state, because it would
+// mean hiding visible (albeit stale) content in the current tree — we prefer to
+// show the stale content, rather than switch to a fallback. But showing a
+// fallback in a new tree is fine, because there's no stale content to
+// prefer instead.
+let shellBoundary: Fiber | null = null;
+
+export function getShellBoundary(): Fiber | null {
+  return shellBoundary;
 }
 
 export function pushPrimaryTreeSuspenseHandler(handler: Fiber): void {
-  const props = handler.pendingProps;
-  const handlerOnStack = suspenseHandlerStackCursor.current;
+  // TODO: Pass as argument
+  const current = handler.alternate;
+  const props: SuspenseProps = handler.pendingProps;
+
+  // Experimental feature: Some Suspense boundaries are marked as having an
+  // undesirable fallback state. These have special behavior where we only
+  // activate the fallback if there's no other boundary on the stack that we can
+  // use instead.
   if (
     enableSuspenseAvoidThisFallback &&
     props.unstable_avoidThisFallback === true &&
-    handlerOnStack !== null &&
-    !shouldAvoidedBoundaryCapture(handler, handlerOnStack, props)
+    // If an avoided boundary is already visible, it behaves identically to
+    // a regular Suspense boundary.
+    (current === null || isCurrentTreeHidden())
   ) {
-    // This boundary should not capture if something suspends. Reuse the
-    // existing handler on the stack.
-    push(suspenseHandlerStackCursor, handlerOnStack, handler);
-  } else {
-    // Push this handler onto the stack.
-    push(suspenseHandlerStackCursor, handler, handler);
+    if (shellBoundary === null) {
+      // We're rendering in the shell. There's no parent Suspense boundary that
+      // can provide a desirable fallback state. We'll use this boundary.
+      push(suspenseHandlerStackCursor, handler, handler);
+
+      // However, because this is not a desirable fallback, the children are
+      // still considered part of the shell. So we intentionally don't assign
+      // to `shellBoundary`.
+    } else {
+      // There's already a parent Suspense boundary that can provide a desirable
+      // fallback state. Prefer that one.
+      const handlerOnStack = suspenseHandlerStackCursor.current;
+      push(suspenseHandlerStackCursor, handlerOnStack, handler);
+    }
+    return;
+  }
+
+  // TODO: If the parent Suspense handler already suspended, there's no reason
+  // to push a nested Suspense handler, because it will get replaced by the
+  // outer fallback, anyway. Consider this as a future optimization.
+  push(suspenseHandlerStackCursor, handler, handler);
+  if (shellBoundary === null) {
+    if (current === null || isCurrentTreeHidden()) {
+      // This boundary is not visible in the current UI.
+      shellBoundary = handler;
+    } else {
+      const prevState: SuspenseState = current.memoizedState;
+      if (prevState !== null) {
+        // This boundary is showing a fallback in the current UI.
+        shellBoundary = handler;
+      }
+    }
   }
 }
 
@@ -115,6 +102,21 @@ export function pushFallbackTreeSuspenseHandler(fiber: Fiber): void {
 export function pushOffscreenSuspenseHandler(fiber: Fiber): void {
   if (fiber.tag === OffscreenComponent) {
     push(suspenseHandlerStackCursor, fiber, fiber);
+    if (shellBoundary !== null) {
+      // A parent boundary is showing a fallback, so we've already rendered
+      // deeper than the shell.
+    } else {
+      const nextState: OffscreenState = fiber.memoizedState;
+      const current = fiber.alternate;
+      if (current !== null) {
+        const prevState: OffscreenState = current.memoizedState;
+        if (prevState !== null) {
+          // This is the first boundary in the stack that's already showing
+          // a fallback. So everything outside is considered the shell.
+          shellBoundary = fiber;
+        }
+      }
+    }
   } else {
     // This is a LegacyHidden component.
     reuseSuspenseHandlerOnStack(fiber);
@@ -131,6 +133,10 @@ export function getSuspenseHandler(): Fiber | null {
 
 export function popSuspenseHandler(fiber: Fiber): void {
   pop(suspenseHandlerStackCursor, fiber);
+  if (shellBoundary === fiber) {
+    // Popping back into the shell.
+    shellBoundary = null;
+  }
 }
 
 // SuspenseList context
diff --git a/packages/react-reconciler/src/ReactFiberThrow.js b/packages/react-reconciler/src/ReactFiberThrow.js
@@ -13,7 +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 type {SuspenseProps, SuspenseState} from './ReactFiberSuspenseComponent';
 
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
 import {
@@ -51,7 +51,10 @@ import {
   enqueueUpdate,
 } from './ReactFiberClassUpdateQueue';
 import {markFailedErrorBoundaryForHotReloading} from './ReactFiberHotReloading';
-import {getSuspenseHandler} from './ReactFiberSuspenseContext';
+import {
+  getShellBoundary,
+  getSuspenseHandler,
+} from './ReactFiberSuspenseContext';
 import {
   renderDidError,
   renderDidSuspendDelayIfPossible,
@@ -368,40 +371,26 @@ function throwException(
           // 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.
+            if (getShellBoundary() === null) {
+              // Suspended in the "shell" of the app. This is an undesirable
+              // loading state. We should avoid committing this tree.
+              renderDidSuspendDelayIfPossible();
             } else {
+              // If we suspended deeper than the shell, we don't need to delay
+              // the commmit. However, we still call renderDidSuspend if this is
+              // a new boundary, to tell the work loop that a new fallback has
+              // appeared during this render.
+              // TODO: Theoretically we should be able to delete this branch.
+              // It's currently used for two things: 1) to throttle the
+              // appearance of successive loading states, and 2) in
+              // SuspenseList, to determine whether the children include any
+              // pending fallbacks. For 1, we should apply throttling to all
+              // retries, not just ones that render an additional fallback. For
+              // 2, we should check subtreeFlags instead. Then we can delete
+              // this branch.
               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();
-                }
+                renderDidSuspend();
               }
             }
           }
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.js b/packages/react-reconciler/src/ReactFiberWorkLoop.js
@@ -276,7 +276,7 @@ import {
 import {schedulePostPaintCallback} from './ReactPostPaintCallback';
 import {
   getSuspenseHandler,
-  isBadSuspenseFallback,
+  getShellBoundary,
 } from './ReactFiberSuspenseContext';
 import {resolveDefaultProps} from './ReactFiberLazyComponent';
 
@@ -1859,19 +1859,11 @@ function handleThrow(root, thrownValue): void {
 }
 
 function shouldAttemptToSuspendUntilDataResolves() {
-  // 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.
-    return true;
-  }
-
   // Check if there are other pending updates that might possibly unblock this
   // component from suspending. This mirrors the check in
   // renderDidSuspendDelayIfPossible. We should attempt to unify them somehow.
+  // TODO: Consider unwinding immediately, using the
+  // SuspendedOnHydration mechanism.
   if (
     includesNonIdleWork(workInProgressRootSkippedLanes) ||
     includesNonIdleWork(workInProgressRootInterleavedUpdatedLanes)
@@ -1884,27 +1876,24 @@ function shouldAttemptToSuspendUntilDataResolves() {
   // TODO: We should be able to remove the equivalent check in
   // finishConcurrentRender, and rely just on this one.
   if (includesOnlyTransitions(workInProgressRootRenderLanes)) {
-    const suspenseHandler = getSuspenseHandler();
-    if (suspenseHandler !== null && suspenseHandler.tag === SuspenseComponent) {
-      const currentSuspenseHandler = suspenseHandler.alternate;
-      const nextProps: SuspenseProps = suspenseHandler.memoizedProps;
-      if (isBadSuspenseFallback(currentSuspenseHandler, nextProps)) {
-        // The nearest Suspense boundary is already showing content. We should
-        // avoid replacing it with a fallback, and instead wait until the
-        // data finishes loading.
-        return true;
-      } else {
-        // This is not a bad fallback condition. We should show a fallback
-        // immediately instead of waiting for the data to resolve. This includes
-        // when suspending inside new trees.
-        return false;
-      }
-    }
+    // If we're rendering inside the "shell" of the app, it's better to suspend
+    // rendering and wait for the data to resolve. Otherwise, we should switch
+    // to a fallback and continue rendering.
+    return getShellBoundary() === null;
+  }
 
-    // During a transition, if there is no Suspense boundary (i.e. suspending in
-    // the "shell" of an application), or if we're inside a hidden tree, then
-    // we should wait until the data finishes loading.
-    return true;
+  const handler = getSuspenseHandler();
+  if (handler === null) {
+    // TODO: We should support suspending in the case where there's no
+    // parent Suspense boundary, even outside a transition. Somehow. Otherwise,
+    // an uncached promise can fall into an infinite loop.
+  } else {
+    if (includesOnlyRetries(workInProgressRootRenderLanes)) {
+      // During a retry, we can suspend rendering if the nearest Suspense boundary
+      // is the boundary of the "shell", because we're guaranteed not to block
+      // any new content from appearing.
+      return handler === getShellBoundary();
+    }
   }
 
   // For all other Lanes besides Transitions and Retries, we should not wait
@@ -1982,6 +1971,8 @@ export function renderDidSuspendDelayIfPossible(): void {
     // (inside this function), since by suspending at the end of the render
     // phase introduces a potential mistake where we suspend lanes that were
     // pinged or updated while we were rendering.
+    // TODO: Consider unwinding immediately, using the
+    // SuspendedOnHydration mechanism.
     markRootSuspended(workInProgressRoot, workInProgressRootRenderLanes);
   }
 }
@@ -2199,6 +2190,10 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
             }
             // The work loop is suspended on data. We should wait for it to
             // resolve before continuing to render.
+            // TODO: Handle the case where the promise resolves synchronously.
+            // Usually this is handled when we instrument the promise to add a
+            // `status` field, but if the promise already has a status, we won't
+            // have added a listener until right here.
             const onResolution = () => {
               ensureRootIsScheduled(root, now());
             };
diff --git a/packages/react-reconciler/src/__tests__/ReactThenable-test.js b/packages/react-reconciler/src/__tests__/ReactThenable-test.js
