[Fizz] Don't double replay elements when it's the postpone point

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;