From bb172e36593ba82c687963812320f93face1af56 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Fri, 13 Oct 2023 00:44:43 -0400
Subject: [PATCH] Don't pop the replay stack if we've already rendered past an
 element

---
 .../ReactDOMFizzStaticBrowser-test.js         | 61 +++++++++++++++++++
 packages/react-server/src/ReactFizzServer.js  | 18 ++++--
 2 files changed, 73 insertions(+), 6 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js
index 9f8bb2fc07fc5..608d668eb1463 100644
--- a/packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js
+++ b/packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js
@@ -1126,4 +1126,65 @@ describe('ReactDOMFizzStaticBrowser', () => {
     // Client rendered
     expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
   });
+
+  // @gate enablePostpone
+  it('can suspend in a replayed component several layers deep', async () => {
+    let prerendering = true;
+    function Postpone() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return 'Hello';
+    }
+
+    let resolve;
+    const promise = new Promise(r => (resolve = r));
+    function Delay({children}) {
+      if (!prerendering) {
+        React.use(promise);
+      }
+      return children;
+    }
+
+    // This wrapper will cause us to do one destructive render past this.
+    function Outer({children}) {
+      return children;
+    }
+
+    function App() {
+      return (
+        <div>
+          <Outer>
+            <Delay>
+              <Suspense fallback="Loading...">
+                <Postpone />
+              </Suspense>
+            </Delay>
+          </Outer>
+        </div>
+      );
+    }
+
+    const prerendered = await ReactDOMFizzStatic.prerender(<App />);
+    expect(prerendered.postponed).not.toBe(null);
+
+    await readIntoContainer(prerendered.prelude);
+
+    prerendering = false;
+
+    const resumedPromise = ReactDOMFizzServer.resume(
+      <App />,
+      JSON.parse(JSON.stringify(prerendered.postponed)),
+    );
+
+    await jest.runAllTimers();
+
+    expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
+
+    await resolve();
+
+    await readIntoContainer(await resumedPromise);
+
+    expect(getVisibleChildren(container)).toEqual(<div>Hello</div>);
+  });
 });
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index 00b1c29843599..66c76ffaebfe7 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -1967,6 +1967,7 @@ function replayElement(
       }
       const childNodes = node[2];
       const childSlots = node[3];
+      const currentNode = task.node;
       task.replay = {nodes: childNodes, slots: childSlots, pendingTasks: 1};
       try {
         renderElement(
@@ -1988,6 +1989,7 @@ function replayElement(
               "The tree doesn't match so React will fallback to client rendering.",
           );
         }
+        task.replay.pendingTasks--;
       } catch (x) {
         if (
           typeof x === 'object' &&
@@ -1995,18 +1997,21 @@ function replayElement(
           (x === SuspenseException || typeof x.then === 'function')
         ) {
           // Suspend
+          if (task.node === currentNode) {
+            // This same element suspended so we need to pop the replay we just added.
+            task.replay = replay;
+          }
           throw x;
         }
+        task.replay.pendingTasks--;
         // Unlike regular render, we don't terminate the siblings if we error
         // during a replay. That's because this component didn't actually error
         // in the original prerender. What's unable to complete is the child
         // replay nodes which might be Suspense boundaries which are able to
         // absorb the error and we can still continue with siblings.
         erroredReplay(request, task.blockedBoundary, x, childNodes, childSlots);
-      } finally {
-        task.replay.pendingTasks--;
-        task.replay = replay;
       }
+      task.replay = replay;
     } else {
       // Let's double check that the component type matches.
       if (type !== REACT_SUSPENSE_TYPE) {
@@ -2370,6 +2375,7 @@ function replayFragment(
             "The tree doesn't match so React will fallback to client rendering.",
         );
       }
+      task.replay.pendingTasks--;
     } catch (x) {
       if (
         typeof x === 'object' &&
@@ -2379,6 +2385,7 @@ function replayFragment(
         // Suspend
         throw x;
       }
+      task.replay.pendingTasks--;
       // Unlike regular render, we don't terminate the siblings if we error
       // during a replay. That's because this component didn't actually error
       // in the original prerender. What's unable to complete is the child
@@ -2386,10 +2393,8 @@ function replayFragment(
       // absorb the error and we can still continue with siblings.
       // This is an error, stash the component stack if it is null.
       erroredReplay(request, task.blockedBoundary, x, childNodes, childSlots);
-    } finally {
-      task.replay.pendingTasks--;
-      task.replay = replay;
     }
+    task.replay = replay;
     // We finished rendering this node, so now we can consume this
     // slot. This must happen after in case we rerender this task.
     replayNodes.splice(j, 1);
@@ -2432,6 +2437,7 @@ function renderChildrenArray(
         // We need to use the non-destructive form so that we can safely pop back
         // up and render the sibling if something suspends.
         const resumeSegmentID = resumeSlots[i];
+        // TODO: If this errors we should still continue with the next sibling.
         if (typeof resumeSegmentID === 'number') {
           resumeNode(request, task, resumeSegmentID, node, i);
           // We finished rendering this node, so now we can consume this