Repository: facebook/react
Pull request: [Fizz] Track postpones in fallbacks

This fixes so that you can postpone in a fallback. This postpones the parent boundary. I track the fallbacks in a separate replay node so that when we resume, we can replay the fallback itself and finish the fallback and then possibly later the content. By doing this we also ensure we don't complete the parent too early since now it has a render task on it.

There is one case that this surfaces that isn't limited to prerender/resume but also render/hydrateRoot. I left todos in the tests for this.

If you postpone in a fallback, and suspend in the content but eventually don't postpone in the content then we should be able to just skip postponing since the content rendered and we no longer need the fallback. This is a bit of a weird edge case though since fallbacks are supposed to be very minimal.

This happens because in both cases the fallback starts rendering early as soon as the content suspends. This also ensures that the parent doesn't complete early by increasing the blocking tasks. Unfortunately, the fallback will irreversibly postpone its parent boundary as soon as it hits a postpone.

When you suspend, the same thing happens but we typically deal with this by doing a "soft" abort on the fallback since we don't need it anymore which unblocks the parent boundary. We can't do that with postpone right now though since it's considered a terminal state.

I think I'll just leave this as is for now since it's an edge case but it's an annoying exception in the model. Makes me feel I haven't quite nailed it just yet.

Changed files:
filename: packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
patch: @@ -6261,6 +6261,59 @@ describe('ReactDOMFizzServer', () => {
     expect(fatalErrors).toEqual(['testing postpone']);
   });
 
+  // @gate enablePostpone
+  it('can postpone in a fallback', async () => {
+    function Postponed({isClient}) {
+      if (!isClient) {
+        React.unstable_postpone('testing postpone');
+      }
+      return 'loading...';
+    }
+
+    const lazyText = React.lazy(async () => {
+      await 0; // causes the fallback to start work
+      return {default: 'Hello'};
+    });
+
+    function App({isClient}) {
+      return (
+        <div>
+          <Suspense fallback="Outer">
+            <Suspense fallback={<Postponed isClient={isClient} />}>
+              {lazyText}
+            </Suspense>
+          </Suspense>
+        </div>
+      );
+    }
+
+    const errors = [];
+
+    await act(() => {
+      const {pipe} = renderToPipeableStream(<App isClient={false} />, {
+        onError(error) {
+          errors.push(error.message);
+        },
+      });
+      pipe(writable);
+    });
+
+    // TODO: This should actually be fully resolved because the value could eventually
+    // resolve on the server even though the fallback couldn't so we should have been
+    // able to render it.
+    expect(getVisibleChildren(container)).toEqual(<div>Outer</div>);
+
+    ReactDOMClient.hydrateRoot(container, <App isClient={true} />, {
+      onRecoverableError(error) {
+        errors.push(error.message);
+      },
+    });
+    await waitForAll([]);
+    // Postponing should not be logged as a recoverable error since it's intentional.
+    expect(errors).toEqual([]);
+    expect(getVisibleChildren(container)).toEqual(<div>Hello</div>);
+  });
+
   it(
     'a transition that flows into a dehydrated boundary should not suspend ' +
       'if the boundary is showing a fallback',
@@ -6830,4 +6883,94 @@ describe('ReactDOMFizzServer', () => {
           ],
     );
   });
+
+  // @gate enablePostpone
+  it('can postpone in fallback', async () => {
+    let prerendering = true;
+    function Postpone() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return 'Hello';
+    }
+
+    let resolve;
+    const promise = new Promise(r => (resolve = r));
+
+    function PostponeAndDelay() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return React.use(promise);
+    }
+
+    const Lazy = React.lazy(async () => {
+      await 0;
+      return {default: Postpone};
+    });
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Outer">
+            <Suspense fallback={<Postpone />}>
+              <PostponeAndDelay /> World
+            </Suspense>
+            <Suspense fallback={<Postpone />}>
+              <Lazy />
+            </Suspense>
+          </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>Outer</div>);
+
+    // Read what we've completed so far
+    await act(() => {
+      resumed.pipe(writable);
+    });
+
+    // Should have now resolved the postponed loading state, but not the promise
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        {'Hello'}
+        {'Hello'}
+      </div>,
+    );
+
+    // Resolve the final promise
+    await act(() => {
+      resolve('Hi');
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        {'Hi'}
+        {' World'}
+        {'Hello'}
+      </div>,
+    );
+  });
 });

filename: packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js
patch: @@ -870,8 +870,6 @@ describe('ReactDOMFizzStaticBrowser', () => {
 
     prerendering = false;
 
-    console.log(JSON.stringify(prerendered.postponed, null, 2));
-
     const resumed = await ReactDOMFizzServer.resume(
       <App />,
       JSON.parse(JSON.stringify(prerendered.postponed)),
@@ -887,4 +885,100 @@ describe('ReactDOMFizzStaticBrowser', () => {
       <div>{['Hello', 'Hello', 'Hello']}</div>,
     );
   });
+
+  // @gate enablePostpone
+  it('can postpone in fallback', async () => {
+    let prerendering = true;
+    function Postpone() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return 'Hello';
+    }
+
+    const Lazy = React.lazy(async () => {
+      await 0;
+      return {default: Postpone};
+    });
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Outer">
+            <Suspense fallback={<Postpone />}>
+              <Postpone /> World
+            </Suspense>
+            <Suspense fallback={<Postpone />}>
+              <Lazy />
+            </Suspense>
+          </Suspense>
+        </div>
+      );
+    }
+
+    const prerendered = await ReactDOMFizzStatic.prerender(<App />);
+    expect(prerendered.postponed).not.toBe(null);
+
+    prerendering = false;
+
+    const resumed = await ReactDOMFizzServer.resume(
+      <App />,
+      JSON.parse(JSON.stringify(prerendered.postponed)),
+    );
+
+    await readIntoContainer(prerendered.prelude);
+
+    expect(getVisibleChildren(container)).toEqual(<div>Outer</div>);
+
+    await readIntoContainer(resumed);
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        {'Hello'}
+        {' World'}
+        {'Hello'}
+      </div>,
+    );
+  });
+
+  // @gate enablePostpone
+  it('can postpone in fallback without postponing the tree', async () => {
+    function Postpone() {
+      React.unstable_postpone();
+    }
+
+    const lazyText = React.lazy(async () => {
+      await 0; // causes the fallback to start work
+      return {default: 'Hello'};
+    });
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Outer">
+            <Suspense fallback={<Postpone />}>{lazyText}</Suspense>
+          </Suspense>
+        </div>
+      );
+    }
+
+    const prerendered = await ReactDOMFizzStatic.prerender(<App />);
+    // TODO: This should actually be null because we should've been able to fully
+    // resolve the render on the server eventually, even though the fallback postponed.
+    // So we should not need to resume.
+    expect(prerendered.postponed).not.toBe(null);
+
+    await readIntoContainer(prerendered.prelude);
+
+    expect(getVisibleChildren(container)).toEqual(<div>Outer</div>);
+
+    const resumed = await ReactDOMFizzServer.resume(
+      <App />,
+      JSON.parse(JSON.stringify(prerendered.postponed)),
+    );
+
+    await readIntoContainer(resumed);
+
+    expect(getVisibleChildren(container)).toEqual(<div>Hello</div>);
+  });
 });

filename: packages/react-server/src/ReactFizzServer.js
patch: @@ -170,8 +170,9 @@ type ResumeSlots =
 type ReplaySuspenseBoundary = [
   string | null /* name */,
   string | number /* key */,
-  Array<ReplayNode> /* keyed children */,
-  ResumeSlots /* resumable slots */,
+  Array<ReplayNode> /* content keyed children */,
+  ResumeSlots /* content resumable slots */,
+  null | ReplayNode /* fallback content */,
   number /* rootSegmentID */,
 ];
 
@@ -208,7 +209,8 @@ type SuspenseBoundary = {
   byteSize: number, // used to determine whether to inline children boundaries.
   fallbackAbortableTasks: Set<Task>, // used to cancel task on the fallback if the boundary completes or gets canceled.
   resources: BoundaryResources,
-  keyPath: Root | KeyNode,
+  trackedContentKeyPath: null | KeyNode, // used to track the path for replay nodes
+  trackedFallbackNode: null | ReplayNode, // used to track the fallback for replay nodes
 };
 
 type RenderTask = {
@@ -549,7 +551,6 @@ function pingTask(request: Request, task: Task): void {
 function createSuspenseBoundary(
   request: Request,
   fallbackAbortableTasks: Set<Task>,
-  keyPath: Root | KeyNode,
 ): SuspenseBoundary {
   return {
     status: PENDING,
@@ -561,7 +562,8 @@ function createSuspenseBoundary(
     fallbackAbortableTasks,
     errorDigest: null,
     resources: createBoundaryResources(),
-    keyPath,
+    trackedContentKeyPath: null,
+    trackedFallbackNode: null,
   };
 }
 
@@ -823,11 +825,10 @@ function renderSuspenseBoundary(
   const content: ReactNodeList = props.children;
 
   const fallbackAbortSet: Set<Task> = new Set();
-  const newBoundary = createSuspenseBoundary(
-    request,
-    fallbackAbortSet,
-    keyPath,
-  );
+  const newBoundary = createSuspenseBoundary(request, fallbackAbortSet);
+  if (request.trackedPostpones !== null) {
+    newBoundary.trackedContentKeyPath = keyPath;
+  }
   const insertionIndex = parentSegment.chunks.length;
   // The children of the boundary segment is actually the fallback.
   const boundarySegment = createPendingSegment(
@@ -930,6 +931,28 @@ function renderSuspenseBoundary(
     task.keyPath = prevKeyPath;
   }
 
+  const fallbackKeyPath = [keyPath[0], 'Suspense Fallback', keyPath[2]];
+  const trackedPostpones = request.trackedPostpones;
+  if (trackedPostpones !== null) {
+    // We create a detached replay node to track any postpones inside the fallback.
+    const fallbackReplayNode: ReplayNode = [
+      fallbackKeyPath[1],
+      fallbackKeyPath[2],
+      ([]: Array<ReplayNode>),
+      null,
+    ];
+    trackedPostpones.workingMap.set(fallbackKeyPath, fallbackReplayNode);
+    if (newBoundary.status === POSTPONED) {
+      // This must exist now.
+      const boundaryReplayNode: ReplaySuspenseBoundary =
+        (trackedPostpones.workingMap.get(keyPath): any);
+      boundaryReplayNode[4] = fallbackReplayNode;
+    } else {
+      // We might not inject it into the postponed tree, unless the content actually
+      // postpones too. We need to keep track of it until that happpens.
+      newBoundary.trackedFallbackNode = fallbackReplayNode;
+    }
+  }
   // 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.
   const suspendedFallbackTask = createRenderTask(
@@ -940,8 +963,7 @@ function renderSuspenseBoundary(
     parentBoundary,
     boundarySegment,
     fallbackAbortSet,
-    // TODO: Should distinguish key path of fallback and primary tasks
-    keyPath,
+    fallbackKeyPath,
     task.formatContext,
     task.legacyContext,
     task.context,
@@ -965,6 +987,8 @@ function replaySuspenseBoundary(
   id: number,
   childNodes: Array<ReplayNode>,
   childSlots: ResumeSlots,
+  fallbackNodes: Array<ReplayNode>,
+  fallbackSlots: ResumeSlots,
 ): void {
   pushBuiltInComponentStackInDEV(task, 'Suspense');
 
@@ -974,13 +998,10 @@ function replaySuspenseBoundary(
   const parentBoundary = task.blockedBoundary;
 
   const content: ReactNodeList = props.children;
+  const fallback: ReactNodeList = props.fallback;
 
   const fallbackAbortSet: Set<Task> = new Set();
-  const resumedBoundary = createSuspenseBoundary(
-    request,
-    fallbackAbortSet,
-    task.keyPath,
-  );
+  const resumedBoundary = createSuspenseBoundary(request, fallbackAbortSet);
   resumedBoundary.parentFlushed = true;
   // We restore the same id of this boundary as was used during prerender.
   resumedBoundary.rootSegmentID = id;
@@ -1003,20 +1024,25 @@ function replaySuspenseBoundary(
     } else {
       renderNode(request, task, content, -1);
     }
-    if (
-      resumedBoundary.pendingTasks === 0 &&
-      resumedBoundary.status === PENDING
-    ) {
-      resumedBoundary.status = COMPLETED;
-      request.completedBoundaries.push(resumedBoundary);
-    }
     if (task.replay.pendingTasks === 1 && task.replay.nodes.length > 0) {
       throw new Error(
         "Couldn't find all resumable slots by key/index during replaying. " +
           "The tree doesn't match so React will fallback to client rendering.",
       );
     }
     task.replay.pendingTasks--;
+    if (
+      resumedBoundary.pendingTasks === 0 &&
+      resumedBoundary.status === PENDING
+    ) {
+      resumedBoundary.status = COMPLETED;
+      request.completedBoundaries.push(resumedBoundary);
+      // This must have been the last segment we were waiting on. This boundary is now complete.
+      // Therefore we won't need the fallback. We early return so that we don't have to create
+      // the fallback.
+      popComponentStackInDEV(task);
+      return;
+    }
   } catch (error) {
     resumedBoundary.status = CLIENT_RENDERED;
     let errorDigest;
@@ -1057,6 +1083,66 @@ function replaySuspenseBoundary(
     task.replay = previousReplaySet;
     task.keyPath = prevKeyPath;
   }
+
+  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,
+    );
+  }
+  if (__DEV__) {
+    suspendedFallbackTask.componentStack = task.componentStack;
+  }
+  // TODO: This should be queued at a separate lower priority queue so that we only work
+  // on preparing fallbacks if we don't have any more main content to task on.
+  request.pingedTasks.push(suspendedFallbackTask);
+
   // TODO: Should this be in the finally?
   popComponentStackInDEV(task);
 }
@@ -2025,9 +2111,11 @@ function replayElement(
         task,
         keyPath,
         props,
-        node[4],
+        node[5],
         node[2],
         node[3],
+        node[4] === null ? [] : node[4][2],
+        node[4] === null ? null : node[4][3],
       );
     }
     // We finished rendering this node, so now we can consume this
@@ -2467,13 +2555,15 @@ function trackPostpone(
     // it before flushing and we know that we can't inline it.
     boundary.rootSegmentID = request.nextSegmentId++;
 
-    const boundaryKeyPath = boundary.keyPath;
+    const boundaryKeyPath = boundary.trackedContentKeyPath;
     if (boundaryKeyPath === null) {
       throw new Error(
         'It should not be possible to postpone at the root. This is a bug in React.',
       );
     }
 
+    const fallbackReplayNode = boundary.trackedFallbackNode;
+
     const children: Array<ReplayNode> = [];
     if (boundaryKeyPath === keyPath && task.childIndex === -1) {
       // Since we postponed directly in the Suspense boundary we can't have written anything
@@ -2485,8 +2575,10 @@ function trackPostpone(
         boundaryKeyPath[2],
         children,
         boundary.rootSegmentID,
+        fallbackReplayNode,
         boundary.rootSegmentID,
       ];
+      trackedPostpones.workingMap.set(boundaryKeyPath, boundaryNode);
       addToReplayParent(boundaryNode, boundaryKeyPath[0], trackedPostpones);
       return;
     } else {
@@ -2498,14 +2590,16 @@ function trackPostpone(
           boundaryKeyPath[2],
           children,
           null,
+          fallbackReplayNode,
           boundary.rootSegmentID,
         ];
         trackedPostpones.workingMap.set(boundaryKeyPath, boundaryNode);
         addToReplayParent(boundaryNode, boundaryKeyPath[0], trackedPostpones);
       } else {
         // Upgrade to ReplaySuspenseBoundary.
-        ((boundaryNode: any): ReplaySuspenseBoundary)[4] =
-          boundary.rootSegmentID;
+        const suspenseBoundary: ReplaySuspenseBoundary = (boundaryNode: any);
+        suspenseBoundary[4] = fallbackReplayNode;
+        suspenseBoundary[5] = boundary.rootSegmentID;
       }
       // Fall through to add the child node.
     }
@@ -2528,13 +2622,19 @@ function trackPostpone(
     if (keyPath === null) {
       trackedPostpones.rootSlots = segment.id;
     } else {
-      const resumableElement: ReplayNode = [
-        keyPath[1],
-        keyPath[2],
-        ([]: Array<ReplayNode>),
-        segment.id,
-      ];
-      addToReplayParent(resumableElement, keyPath[0], trackedPostpones);
+      const workingMap = trackedPostpones.workingMap;
+      let resumableNode = workingMap.get(keyPath);
+      if (resumableNode === undefined) {
+        resumableNode = [
+          keyPath[1],
+          keyPath[2],
+          ([]: Array<ReplayNode>),
+          segment.id,
+        ];
+        addToReplayParent(resumableNode, keyPath[0], trackedPostpones);
+      } else {
+        resumableNode[3] = segment.id;
+      }
     }
   } else {
     let slots;
@@ -2963,11 +3063,7 @@ function abortRemainingSuspenseBoundary(
   error: mixed,
   errorDigest: ?string,
 ): void {
-  const resumedBoundary = createSuspenseBoundary(
-    request,
-    new Set(),
-    null, // The keyPath doesn't matter at this point so we don't bother rebuilding it.
-  );
+  const resumedBoundary = createSuspenseBoundary(request, new Set());
   resumedBoundary.parentFlushed = true;
   // We restore the same id of this boundary as was used during prerender.
   resumedBoundary.rootSegmentID = rootSegmentID;
@@ -3017,7 +3113,7 @@ function abortRemainingReplayNodes(
       );
     } else {
       const boundaryNode: ReplaySuspenseBoundary = node;
-      const rootSegmentID = boundaryNode[4];
+      const rootSegmentID = boundaryNode[5];
       abortRemainingSuspenseBoundary(
         request,
         rootSegmentID,
@@ -3835,9 +3931,7 @@ function flushCompletedQueues(
             destination,
             request.resumableState,
             request.renderState,
-            request.allPendingTasks === 0 &&
-              (request.trackedPostpones === null ||
-                request.trackedPostpones.workingMap.size === 0),
+            request.allPendingTasks === 0 && request.trackedPostpones === null,
           );
         }
 
@@ -3932,13 +4026,7 @@ function flushCompletedQueues(
       if (enableFloat) {
         // We write the trailing tags but only if don't have any data to resume.
         // If we need to resume we'll write the postamble in the resume instead.
-        if (
-          !enablePostpone ||
-          request.trackedPostpones === null ||
-          // We check the working map instead of the root because the root could've
-          // been mutated at this point if it was passed straight through to resume().
-          request.trackedPostpones.workingMap.size === 0
-        ) {
+        if (!enablePostpone || request.trackedPostpones === null) {
           writePostamble(destination, request.resumableState);
         }
       }
@@ -4090,6 +4178,8 @@ export function getPostponedState(request: Request): null | PostponedState {
     (trackedPostpones.rootNodes.length === 0 &&
       trackedPostpones.rootSlots === null)
   ) {
+    // Reset. Let the flushing behave as if we completed the whole document.
+    request.trackedPostpones = null;
     return null;
   }
   return {