Retrying a dehydrated boundary pings at the earliest forced time

This might quickly become an already expired time.

Author: sebmarkbage

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -171,7 +171,7 @@ import {
 import {
   markSpawnedWork,
   requestCurrentTime,
-  retryTimedOutBoundary,
+  retryDehydratedSuspenseBoundary,
   scheduleWork,
 } from './ReactFiberWorkLoop';
 
@@ -2081,7 +2081,7 @@ function updateDehydratedSuspenseComponent(
     // Register a callback to retry this boundary once the server has sent the result.
     registerSuspenseInstanceRetry(
       suspenseInstance,
-      retryTimedOutBoundary.bind(null, current),
+      retryDehydratedSuspenseBoundary.bind(null, current),
     );
     return null;
   } else {

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -16,6 +16,7 @@ import type {
 } from './SchedulerWithReactIntegration';
 import type {Interaction} from 'scheduler/src/Tracing';
 import type {SuspenseConfig} from './ReactFiberSuspenseConfig';
+import type {SuspenseState} from './ReactFiberSuspenseComponent';
 
 import {
   warnAboutDeprecatedLifecycles,
@@ -2183,18 +2184,23 @@ export function pingSuspendedRoot(
   scheduleCallbackForRoot(root, priorityLevel, suspendedTime);
 }
 
-export function retryTimedOutBoundary(boundaryFiber: Fiber) {
+function retryTimedOutBoundary(
+  boundaryFiber: Fiber,
+  retryTime: ExpirationTime,
+) {
   // The boundary fiber (a Suspense component or SuspenseList component)
   // previously was rendered in its fallback state. One of the promises that
   // suspended it has resolved, which means at least part of the tree was
   // likely unblocked. Try rendering again, at a new expiration time.
   const currentTime = requestCurrentTime();
-  const suspenseConfig = null; // Retries don't carry over the already committed update.
-  const retryTime = computeExpirationForFiber(
-    currentTime,
-    boundaryFiber,
-    suspenseConfig,
-  );
+  if (retryTime === Never) {
+    const suspenseConfig = null; // Retries don't carry over the already committed update.
+    retryTime = computeExpirationForFiber(
+      currentTime,
+      boundaryFiber,
+      suspenseConfig,
+    );
+  }
   // TODO: Special case idle priority?
   const priorityLevel = inferPriorityFromExpirationTime(currentTime, retryTime);
   const root = markUpdateTimeFromFiberToRoot(boundaryFiber, retryTime);
@@ -2203,12 +2209,26 @@ export function retryTimedOutBoundary(boundaryFiber: Fiber) {
   }
 }
 
+export function retryDehydratedSuspenseBoundary(boundaryFiber: Fiber) {
+  const suspenseState: null | SuspenseState = boundaryFiber.memoizedState;
+  let retryTime = Never;
+  if (suspenseState !== null) {
+    retryTime = suspenseState.retryTime;
+  }
+  retryTimedOutBoundary(boundaryFiber, retryTime);
+}
+
 export function resolveRetryThenable(boundaryFiber: Fiber, thenable: Thenable) {
+  let retryTime = Never; // Default
   let retryCache: WeakSet<Thenable> | Set<Thenable> | null;
   if (enableSuspenseServerRenderer) {
     switch (boundaryFiber.tag) {
       case SuspenseComponent:
         retryCache = boundaryFiber.stateNode;
+        const suspenseState: null | SuspenseState = boundaryFiber.memoizedState;
+        if (suspenseState !== null) {
+          retryTime = suspenseState.retryTime;
+        }
         break;
       default:
         invariant(
@@ -2227,7 +2247,7 @@ export function resolveRetryThenable(boundaryFiber: Fiber, thenable: Thenable) {
     retryCache.delete(thenable);
   }
 
-  retryTimedOutBoundary(boundaryFiber);
+  retryTimedOutBoundary(boundaryFiber, retryTime);
 }
 
 // Computes the next Just Noticeable Difference (JND) boundary.