[Fizz] Don't double replay elements when it's the postpone point (#27440)

The `resumeElement` function wasn't actually doing the correct thing
because it was resuming the element itself but what the child slot means
is that we're supposed to resume in the direct child of the element.
This is difficult to check for since it's all the possible branches that
the element can render into, so instead we just check this in
renderNode. It means the hottest path always checks the task which is
unfortunate.

And also, resuming using the correct nextSegmentId.

Fixes two bugs surfaced by this test.

---------

Co-authored-by: Josh Story <[email protected]>

Author: sebmarkbage

packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js

@@ -6988,4 +6988,100 @@ describe('ReactDOMFizzServer', () => {
       </div>,
     );
   });
+
+  // @gate enablePostpone
+  it('can discover new suspense boundaries in the resume', async () => {
+    let prerendering = true;
+    let resolveA;
+    const promiseA = new Promise(r => (resolveA = r));
+    let resolveB;
+    const promiseB = new Promise(r => (resolveB = r));
+
+    function WaitA() {
+      return React.use(promiseA);
+    }
+    function WaitB() {
+      return React.use(promiseB);
+    }
+    function Postpone() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return (
+        <span>
+          <Suspense fallback="Loading again...">
+            <WaitA />
+          </Suspense>
+          <WaitB />
+        </span>
+      );
+    }
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <p>
+              <Postpone />
+            </p>
+          </Suspense>
+        </div>
+      );
+    }
+
+    const prerendered = await ReactDOMFizzStatic.prerenderToNodeStream(<App />);
+    expect(prerendered.postponed).not.toBe(null);
+
+    prerendering = false;
+
+    // Create a separate stream so it doesn't close the writable. I.e. simple concat.
+    const preludeWritable = new Stream.PassThrough();
+    preludeWritable.setEncoding('utf8');
+    preludeWritable.on('data', chunk => {
+      writable.write(chunk);
+    });
+
+    await act(() => {
+      prerendered.prelude.pipe(preludeWritable);
+    });
+
+    const resumed = await ReactDOMFizzServer.resumeToPipeableStream(
+      <App />,
+      JSON.parse(JSON.stringify(prerendered.postponed)),
+    );
+
+    expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
+
+    // Read what we've completed so far
+    await act(() => {
+      resumed.pipe(writable);
+    });
+
+    // Still blocked
+    expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
+
+    // Resolve the first promise, this unblocks the inner boundary
+    await act(() => {
+      resolveA('Hello');
+    });
+
+    // Still blocked
+    expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
+
+    // Resolve the second promise, this unblocks the outer boundary
+    await act(() => {
+      resolveB('World');
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <p>
+          <span>
+            {'Hello'}
+            {'World'}
+          </span>
+        </p>
+      </div>,
+    );
+  });
 });

packages/react-server/src/ReactFizzServer.js

@@ -487,7 +487,7 @@ export function resumeRequest(
     progressiveChunkSize: postponedState.progressiveChunkSize,
     status: OPEN,
     fatalError: null,
-    nextSegmentId: 0,
+    nextSegmentId: postponedState.nextSegmentId,
     allPendingTasks: 0,
     pendingRootTasks: 0,
     completedRootSegment: null,
@@ -1019,11 +1019,7 @@ function replaySuspenseBoundary(
   }
   try {
     // We use the safe form because we don't handle suspending here. Only error handling.
-    if (typeof childSlots === 'number') {
-      resumeNode(request, task, childSlots, content, -1);
-    } else {
-      renderNode(request, task, content, -1);
-    }
+    renderNode(request, task, content, -1);
     if (task.replay.pendingTasks === 1 && task.replay.nodes.length > 0) {
       throw new Error(
         "Couldn't find all resumable slots by key/index during replaying. " +
@@ -1086,56 +1082,27 @@ function replaySuspenseBoundary(
 
   const fallbackKeyPath = [keyPath[0], 'Suspense Fallback', keyPath[2]];
 
-  let suspendedFallbackTask;
   // We create suspended task for the fallback because we don't want to actually work
   // on it yet in case we finish the main content, so we queue for later.
-  if (typeof fallbackSlots === 'number') {
-    // Resuming directly in the fallback.
-    const resumedSegment = createPendingSegment(
-      request,
-      0,
-      null,
-      task.formatContext,
-      false,
-      false,
-    );
-    resumedSegment.id = fallbackSlots;
-    resumedSegment.parentFlushed = true;
-    suspendedFallbackTask = createRenderTask(
-      request,
-      null,
-      fallback,
-      -1,
-      parentBoundary,
-      resumedSegment,
-      fallbackAbortSet,
-      fallbackKeyPath,
-      task.formatContext,
-      task.legacyContext,
-      task.context,
-      task.treeContext,
-    );
-  } else {
-    const fallbackReplay = {
-      nodes: fallbackNodes,
-      slots: fallbackSlots,
-      pendingTasks: 0,
-    };
-    suspendedFallbackTask = createReplayTask(
-      request,
-      null,
-      fallbackReplay,
-      fallback,
-      -1,
-      parentBoundary,
-      fallbackAbortSet,
-      fallbackKeyPath,
-      task.formatContext,
-      task.legacyContext,
-      task.context,
-      task.treeContext,
-    );
-  }
+  const fallbackReplay = {
+    nodes: fallbackNodes,
+    slots: fallbackSlots,
+    pendingTasks: 0,
+  };
+  const suspendedFallbackTask = createReplayTask(
+    request,
+    null,
+    fallbackReplay,
+    fallback,
+    -1,
+    parentBoundary,
+    fallbackAbortSet,
+    fallbackKeyPath,
+    task.formatContext,
+    task.legacyContext,
+    task.context,
+    task.treeContext,
+  );
   if (__DEV__) {
     suspendedFallbackTask.componentStack = task.componentStack;
   }
@@ -1965,50 +1932,6 @@ function resumeNode(
   }
 }
 
-function resumeElement(
-  request: Request,
-  task: ReplayTask,
-  keyPath: KeyNode,
-  segmentId: number,
-  prevThenableState: ThenableState | null,
-  type: any,
-  props: Object,
-  ref: any,
-): void {
-  const prevReplay = task.replay;
-  const blockedBoundary = task.blockedBoundary;
-  const resumedSegment = createPendingSegment(
-    request,
-    0,
-    null,
-    task.formatContext,
-    false,
-    false,
-  );
-  resumedSegment.id = segmentId;
-  resumedSegment.parentFlushed = true;
-  try {
-    // Convert the current ReplayTask to a RenderTask.
-    const renderTask: RenderTask = (task: any);
-    renderTask.replay = null;
-    renderTask.blockedSegment = resumedSegment;
-    renderElement(request, task, keyPath, prevThenableState, type, props, ref);
-    resumedSegment.status = COMPLETED;
-    if (blockedBoundary === null) {
-      request.completedRootSegment = resumedSegment;
-    } else {
-      queueCompletedSegment(blockedBoundary, resumedSegment);
-      if (blockedBoundary.parentFlushed) {
-        request.partialBoundaries.push(blockedBoundary);
-      }
-    }
-  } finally {
-    // Restore to a ReplayTask.
-    task.replay = prevReplay;
-    task.blockedSegment = null;
-  }
-}
-
 function replayElement(
   request: Request,
   task: ReplayTask,
@@ -2045,29 +1968,15 @@ function replayElement(
       const childSlots = node[3];
       task.replay = {nodes: childNodes, slots: childSlots, pendingTasks: 1};
       try {
-        if (typeof childSlots === 'number') {
-          // Matched a resumable element.
-          resumeElement(
-            request,
-            task,
-            keyPath,
-            childSlots,
-            prevThenableState,
-            type,
-            props,
-            ref,
-          );
-        } else {
-          renderElement(
-            request,
-            task,
-            keyPath,
-            prevThenableState,
-            type,
-            props,
-            ref,
-          );
-        }
+        renderElement(
+          request,
+          task,
+          keyPath,
+          prevThenableState,
+          type,
+          props,
+          ref,
+        );
         if (
           task.replay.pendingTasks === 1 &&
           task.replay.nodes.length > 0
@@ -2215,6 +2124,12 @@ function renderNodeDestructiveImpl(
   node: ReactNodeList,
   childIndex: number,
 ): void {
+  if (task.replay !== null && typeof task.replay.slots === 'number') {
+    // TODO: Figure out a cheaper place than this hot path to do this check.
+    const resumeSegmentID = task.replay.slots;
+    resumeNode(request, task, resumeSegmentID, node, childIndex);
+    return;
+  }
   // Stash the node we're working on. We'll pick up from this task in case
   // something suspends.
   task.node = node;