From de68d2f4a2403ad1ef46a3036ddc1f9080640588 Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Wed, 4 Dec 2024 07:58:43 -0800
Subject: [PATCH] Register Suspense retry handlers in commit phase (#31667)

To avoid GC pressure and accidentally hanging onto old trees Suspense
boundary retries are now implemented in the commit phase. I used the
Callback flag which was previously only used to schedule callbacks for
Class components. This isn't quite semantically equivalent but it's
unused and seemingly compatible.
---
 .../src/client/ReactFiberConfigDOM.js         | 34 ++++++++++++-------
 .../src/ReactFiberBeginWork.js                |  9 ++---
 .../src/ReactFiberCommitWork.js               | 19 +++++++++++
 3 files changed, 43 insertions(+), 19 deletions(-)

diff --git a/packages/react-dom-bindings/src/client/ReactFiberConfigDOM.js b/packages/react-dom-bindings/src/client/ReactFiberConfigDOM.js
index bac5ec8fb601b..0933794c68cb6 100644
--- a/packages/react-dom-bindings/src/client/ReactFiberConfigDOM.js
+++ b/packages/react-dom-bindings/src/client/ReactFiberConfigDOM.js
@@ -1310,20 +1310,28 @@ export function registerSuspenseInstanceRetry(
   callback: () => void,
 ) {
   const ownerDocument = instance.ownerDocument;
-  if (ownerDocument.readyState !== DOCUMENT_READY_STATE_COMPLETE) {
-    ownerDocument.addEventListener(
-      'DOMContentLoaded',
-      () => {
-        if (instance.data === SUSPENSE_PENDING_START_DATA) {
-          callback();
-        }
-      },
-      {
-        once: true,
-      },
-    );
+  if (
+    // The Fizz runtime must have put this boundary into client render or complete
+    // state after the render finished but before it committed. We need to call the
+    // callback now rather than wait
+    instance.data !== SUSPENSE_PENDING_START_DATA ||
+    // The boundary is still in pending status but the document has finished loading
+    // before we could register the event handler that would have scheduled the retry
+    // on load so we call teh callback now.
+    ownerDocument.readyState === DOCUMENT_READY_STATE_COMPLETE
+  ) {
+    callback();
+  } else {
+    // We're still in pending status and the document is still loading so we attach
+    // a listener to the document load even and expose the retry on the instance for
+    // the Fizz runtime to trigger if it ends up resolving this boundary
+    const listener = () => {
+      callback();
+      ownerDocument.removeEventListener('DOMContentLoaded', listener);
+    };
+    ownerDocument.addEventListener('DOMContentLoaded', listener);
+    instance._reactRetry = listener;
   }
-  instance._reactRetry = callback;
 }
 
 export function canHydrateFormStateMarker(
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index 0a5457d20c33b..df222199b2012 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.js
@@ -79,6 +79,7 @@ import {
   PerformedWork,
   Placement,
   Hydrating,
+  Callback,
   ContentReset,
   DidCapture,
   Update,
@@ -166,7 +167,6 @@ import {
   isSuspenseInstancePending,
   isSuspenseInstanceFallback,
   getSuspenseInstanceFallbackErrorDetails,
-  registerSuspenseInstanceRetry,
   supportsHydration,
   supportsResources,
   supportsSingletons,
@@ -256,7 +256,6 @@ import {
   isFunctionClassComponent,
 } from './ReactFiber';
 import {
-  retryDehydratedSuspenseBoundary,
   scheduleUpdateOnFiber,
   renderDidSuspendDelayIfPossible,
   markSkippedUpdateLanes,
@@ -2816,12 +2815,10 @@ function updateDehydratedSuspenseComponent(
       // on the client than if we just leave it alone. If the server times out or errors
       // these should update this boundary to the permanent Fallback state instead.
       // Mark it as having captured (i.e. suspended).
-      workInProgress.flags |= DidCapture;
+      // Also Mark it as requiring retry.
+      workInProgress.flags |= DidCapture | Callback;
       // Leave the child in place. I.e. the dehydrated fragment.
       workInProgress.child = current.child;
-      // Register a callback to retry this boundary once the server has sent the result.
-      const retry = retryDehydratedSuspenseBoundary.bind(null, current);
-      registerSuspenseInstanceRetry(suspenseInstance, retry);
       return null;
     } else {
       // This is the first attempt.
diff --git a/packages/react-reconciler/src/ReactFiberCommitWork.js b/packages/react-reconciler/src/ReactFiberCommitWork.js
index 5c64a0407048e..7887be44e46c0 100644
--- a/packages/react-reconciler/src/ReactFiberCommitWork.js
+++ b/packages/react-reconciler/src/ReactFiberCommitWork.js
@@ -142,6 +142,7 @@ import {
   suspendInstance,
   suspendResource,
   resetFormInstance,
+  registerSuspenseInstanceRetry,
 } from './ReactFiberConfig';
 import {
   captureCommitPhaseError,
@@ -154,6 +155,7 @@ import {
   addMarkerProgressCallbackToPendingTransition,
   addMarkerIncompleteCallbackToPendingTransition,
   addMarkerCompleteCallbackToPendingTransition,
+  retryDehydratedSuspenseBoundary,
 } from './ReactFiberWorkLoop';
 import {
   HasEffect as HookHasEffect,
@@ -526,6 +528,23 @@ function commitLayoutEffectOnFiber(
       if (flags & Update) {
         commitSuspenseHydrationCallbacks(finishedRoot, finishedWork);
       }
+      if (flags & Callback) {
+        // This Boundary is in fallback and has a dehydrated Suspense instance.
+        // We could in theory assume the dehydrated state but we recheck it for
+        // certainty.
+        const finishedState: SuspenseState | null = finishedWork.memoizedState;
+        if (finishedState !== null) {
+          const dehydrated = finishedState.dehydrated;
+          if (dehydrated !== null) {
+            // Register a callback to retry this boundary once the server has sent the result.
+            const retry = retryDehydratedSuspenseBoundary.bind(
+              null,
+              finishedWork,
+            );
+            registerSuspenseInstanceRetry(dehydrated, retry);
+          }
+        }
+      }
       break;
     }
     case OffscreenComponent: {