[Fizz] Postponing in the shell (#27569)

sebmarkbage · web-flow · commit 05fbd1aab036 · 2023-10-23T15:53:59.000-04:00
When we postpone a prerender in the shell, we should just leave an empty
prelude and resume from the root. While preserving any options passed
in.

Since we haven't flushed anything we can't assume we've already emitted
html/body tags or any resources tracked in the resumable state. This
introduces a resetResumableState function to reset anything we didn't
flush.

This is a bit hacky. Ideally, we probably shouldn't have tracked it as
already happened until it flushed or something like that.

Basically, it's like restarting the prerender with the same options and
then immediately aborting. When we add the preload headers, we'd track
those as preload() being emitted after the reset and so they get readded
to the resumable state in that case.
diff --git a/packages/react-dom-bindings/src/server/ReactFizzConfigDOM.js b/packages/react-dom-bindings/src/server/ReactFizzConfigDOM.js
@@ -576,6 +576,29 @@ export function createResumableState(
   };
 }
 
+export function resetResumableState(
+  resumableState: ResumableState,
+  renderState: RenderState,
+): void {
+  // Resets the resumable state based on what didn't manage to fully flush in the render state.
+  // This currently assumes nothing was flushed.
+  resumableState.nextFormID = 0;
+  resumableState.hasBody = false;
+  resumableState.hasHtml = false;
+  resumableState.unknownResources = {};
+  resumableState.dnsResources = {};
+  resumableState.connectResources = {
+    default: {},
+    anonymous: {},
+    credentials: {},
+  };
+  resumableState.imageResources = {};
+  resumableState.styleResources = {};
+  resumableState.scriptResources = {};
+  resumableState.moduleUnknownResources = {};
+  resumableState.moduleScriptResources = {};
+}
+
 // Constants for the insertion mode we're currently writing in. We don't encode all HTML5 insertion
 // modes. We only include the variants as they matter for the sake of our purposes.
 // We don't actually provide the namespace therefore we use constants instead of the string.
diff --git a/packages/react-dom-bindings/src/server/ReactFizzConfigDOMLegacy.js b/packages/react-dom-bindings/src/server/ReactFizzConfigDOMLegacy.js
@@ -158,6 +158,7 @@ export {
   hoistResources,
   setCurrentlyRenderingBoundaryResourcesTarget,
   prepareHostDispatcher,
+  resetResumableState,
 } from './ReactFizzConfigDOM';
 
 import escapeTextForBrowser from './escapeTextForBrowser';
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js
@@ -1187,4 +1187,133 @@ describe('ReactDOMFizzStaticBrowser', () => {
 
     expect(getVisibleChildren(container)).toEqual(<div>Hello</div>);
   });
+
+  // @gate enablePostpone
+  it('emits an empty prelude and resumes at the root if we postpone in the shell', async () => {
+    let prerendering = true;
+    function Postpone() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return 'Hello';
+    }
+
+    function App() {
+      return (
+        <html lang="en">
+          <body>
+            <link rel="stylesheet" href="my-style" precedence="high" />
+            <Postpone />
+          </body>
+        </html>
+      );
+    }
+
+    const prerendered = await ReactDOMFizzStatic.prerender(<App />);
+    expect(prerendered.postponed).not.toBe(null);
+
+    prerendering = false;
+
+    expect(await readContent(prerendered.prelude)).toBe('');
+
+    const content = await ReactDOMFizzServer.resume(
+      <App />,
+      JSON.parse(JSON.stringify(prerendered.postponed)),
+    );
+
+    expect(await readContent(content)).toBe(
+      '<!DOCTYPE html><html lang="en"><head>' +
+        '<link rel="stylesheet" href="my-style" data-precedence="high"/>' +
+        '</head><body>Hello</body></html>',
+    );
+  });
+
+  // @gate enablePostpone
+  it('emits an empty prelude if we have not rendered html or head tags yet', async () => {
+    let prerendering = true;
+    function Postpone() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return (
+        <html lang="en">
+          <body>Hello</body>
+        </html>
+      );
+    }
+
+    function App() {
+      return (
+        <>
+          <link rel="stylesheet" href="my-style" precedence="high" />
+          <Postpone />
+        </>
+      );
+    }
+
+    const prerendered = await ReactDOMFizzStatic.prerender(<App />);
+    expect(prerendered.postponed).not.toBe(null);
+
+    prerendering = false;
+
+    expect(await readContent(prerendered.prelude)).toBe('');
+
+    const content = await ReactDOMFizzServer.resume(
+      <App />,
+      JSON.parse(JSON.stringify(prerendered.postponed)),
+    );
+
+    expect(await readContent(content)).toBe(
+      '<!DOCTYPE html><html lang="en"><head>' +
+        '<link rel="stylesheet" href="my-style" data-precedence="high"/>' +
+        '</head><body>Hello</body></html>',
+    );
+  });
+
+  // @gate enablePostpone
+  it('emits an empty prelude if a postpone in a promise in the shell', 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 (
+        <html>
+          <link rel="stylesheet" href="my-style" precedence="high" />
+          <body>
+            <div>
+              <Lazy />
+            </div>
+          </body>
+        </html>
+      );
+    }
+
+    const prerendered = await ReactDOMFizzStatic.prerender(<App />);
+    expect(prerendered.postponed).not.toBe(null);
+
+    prerendering = false;
+
+    expect(await readContent(prerendered.prelude)).toBe('');
+
+    const content = await ReactDOMFizzServer.resume(
+      <App />,
+      JSON.parse(JSON.stringify(prerendered.postponed)),
+    );
+
+    expect(await readContent(content)).toBe(
+      '<!DOCTYPE html><html><head>' +
+        '<link rel="stylesheet" href="my-style" data-precedence="high"/>' +
+        '</head><body><div>Hello</div></body></html>',
+    );
+  });
 });
diff --git a/packages/react-noop-renderer/src/ReactNoopServer.js b/packages/react-noop-renderer/src/ReactNoopServer.js
@@ -94,6 +94,8 @@ const ReactNoopServer = ReactFizzServer({
     return null;
   },
 
+  resetResumableState(): void {},
+
   pushTextInstance(
     target: Array<Uint8Array>,
     text: string,
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
@@ -76,6 +76,7 @@ import {
   requestStorage,
   pushFormStateMarkerIsMatching,
   pushFormStateMarkerIsNotMatching,
+  resetResumableState,
 } from './ReactFizzConfig';
 import {
   constructClassInstance,
@@ -505,6 +506,39 @@ export function resumeRequest(
     onFatalError: onFatalError === undefined ? noop : onFatalError,
     formState: null,
   };
+  if (typeof postponedState.replaySlots === 'number') {
+    const resumedId = postponedState.replaySlots;
+    // We have a resume slot at the very root. This is effectively just a full rerender.
+    const rootSegment = createPendingSegment(
+      request,
+      0,
+      null,
+      postponedState.rootFormatContext,
+      // Root segments are never embedded in Text on either edge
+      false,
+      false,
+    );
+    rootSegment.id = resumedId;
+    // There is no parent so conceptually, we're unblocked to flush this segment.
+    rootSegment.parentFlushed = true;
+    const rootTask = createRenderTask(
+      request,
+      null,
+      children,
+      -1,
+      null,
+      rootSegment,
+      abortSet,
+      null,
+      postponedState.rootFormatContext,
+      emptyContextObject,
+      rootContextSnapshot,
+      emptyTreeContext,
+    );
+    pingedTasks.push(rootTask);
+    return request;
+  }
+
   const replay: ReplaySet = {
     nodes: postponedState.replayNodes,
     slots: postponedState.replaySlots,
@@ -2477,6 +2511,17 @@ function trackPostpone(
 
   const keyPath = task.keyPath;
   const boundary = task.blockedBoundary;
+
+  if (boundary === null) {
+    segment.id = request.nextSegmentId++;
+    trackedPostpones.rootSlots = segment.id;
+    if (request.completedRootSegment !== null) {
+      // Postpone the root if this was a deeper segment.
+      request.completedRootSegment.status = POSTPONED;
+    }
+    return;
+  }
+
   if (boundary !== null && boundary.status === PENDING) {
     boundary.status = POSTPONED;
     // We need to eagerly assign it an ID because we'll need to refer to
@@ -2835,7 +2880,7 @@ function renderNode(
           enablePostpone &&
           request.trackedPostpones !== null &&
           x.$$typeof === REACT_POSTPONE_TYPE &&
-          task.blockedBoundary !== null // TODO: Support holes in the shell
+          task.blockedBoundary !== null // bubble if we're postponing in the shell
         ) {
           // If we're tracking postpones, we inject a hole here and continue rendering
           // sibling. Similar to suspending. If we're not tracking, we treat it more like
@@ -3376,8 +3421,7 @@ function retryRenderTask(
       } else if (
         enablePostpone &&
         request.trackedPostpones !== null &&
-        x.$$typeof === REACT_POSTPONE_TYPE &&
-        task.blockedBoundary !== null // TODO: Support holes in the shell
+        x.$$typeof === REACT_POSTPONE_TYPE
       ) {
         // If we're tracking postpones, we mark this segment as postponed and finish
         // the task without filling it in. If we're not tracking, we treat it more like
@@ -3870,7 +3914,10 @@ function flushCompletedQueues(
     let i;
     const completedRootSegment = request.completedRootSegment;
     if (completedRootSegment !== null) {
-      if (request.pendingRootTasks === 0) {
+      if (completedRootSegment.status === POSTPONED) {
+        // We postponed the root, so we write nothing.
+        return;
+      } else if (request.pendingRootTasks === 0) {
         if (enableFloat) {
           writePreamble(
             destination,
@@ -4138,6 +4185,13 @@ export function getPostponedState(request: Request): null | PostponedState {
     request.trackedPostpones = null;
     return null;
   }
+  if (
+    request.completedRootSegment !== null &&
+    request.completedRootSegment.status === POSTPONED
+  ) {
+    // We postponed the root so we didn't flush anything.
+    resetResumableState(request.resumableState, request.renderState);
+  }
   return {
     nextSegmentId: request.nextSegmentId,
     rootFormatContext: request.rootFormatContext,
diff --git a/packages/react-server/src/forks/ReactFizzConfig.custom.js b/packages/react-server/src/forks/ReactFizzConfig.custom.js
@@ -39,6 +39,7 @@ export const isPrimaryRenderer = false;
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request> = (null: any);
 
+export const resetResumableState = $$$config.resetResumableState;
 export const getChildFormatContext = $$$config.getChildFormatContext;
 export const makeId = $$$config.makeId;
 export const pushTextInstance = $$$config.pushTextInstance;
