[Fizz] Don't pop the replay stack if we've already rendered past an e…

…lement (#27513)

This is the same problem as we had with keyPath before where if the
element itself suspends, we have to restore the replay node to what it
was before, however, if something below the element suspends we
shouldn't pop it because that will pop it back up the stack.

Instead of passing replay as an argument to every renderElement
function, I use a hack to compare if the node is still the same as the
one we tried to render, then that means we haven't stepped down into the
child yet. Maybe this is not quite correct because in theory you could
have a recursive node that just renders itself over and over until some
context bails out.

This solves an issue where if you suspended in an element it would retry
trying to replay from that element but using the postponed state from
the root.

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>);
+  });
 });

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,25 +1989,29 @@ function replayElement(
               "The tree doesn't match so React will fallback to client rendering.",
           );
         }
+        task.replay.pendingTasks--;
       } catch (x) {
         if (
           typeof x === 'object' &&
           x !== null &&
           (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,17 +2385,16 @@ 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
       // replay nodes which might be Suspense boundaries which are able to
       // 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