From 1219d57fc9fcbf44c873c0b10e5acbd31f613c15 Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Thu, 1 Feb 2024 07:14:08 -0800
Subject: [PATCH] [Fizz] Support aborting with Postpone (#28183)

Semantically if you make your reason for aborting a Postpone instance
the render should not hit the error pathways but should instead follow
the postpone pathways. It's awkward today to actually get your hands on
a Postpone instance because you have to catch the throw from postpone
and then pass that into `abort()` or `AbortController.abort()`
(depending on the renderer API you are using)

This change makes it so that in most circumstances if you abort with a
postpone the `onPostpone` handler will be called and the Suspense
boundaries still pending will be put into client render mode with the
appropriate postpone digest to avoid trigger recoverable error pathways
on the client.

Similar to postponing in the shell during a resume or render however if
you abort before the shell is complete in a resume or render we will
fatally error. The fatal error is contextualized by React to avoid
passing the postpone object itself to the `onError` and related options.
---
 .../src/__tests__/ReactDOMFizzServer-test.js  | 447 ++++++++++++++++++
 packages/react-server/src/ReactFizzServer.js  |  49 +-
 scripts/error-codes/codes.json                |   3 +-
 3 files changed, 494 insertions(+), 5 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
index d0dcb2f714743..da250d55558fe 100644
--- a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
+++ b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
@@ -7498,4 +7498,451 @@ describe('ReactDOMFizzServer', () => {
       </div>,
     );
   });
+
+  // @gate enablePostpone
+  it('does not call onError when you abort with a postpone instance during prerender', async () => {
+    const promise = new Promise(r => {});
+
+    function Wait() {
+      return React.use(promise);
+    }
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <p>
+              <span>
+                <Suspense fallback="Loading again...">
+                  <Wait />
+                </Suspense>
+              </span>
+            </p>
+            <p>
+              <span>
+                <Suspense fallback="Loading again too...">
+                  <Wait />
+                </Suspense>
+              </span>
+            </p>
+          </Suspense>
+        </div>
+      );
+    }
+
+    let postponeInstance;
+    try {
+      React.unstable_postpone('manufactured');
+    } catch (p) {
+      postponeInstance = p;
+    }
+
+    const controller = new AbortController();
+    const signal = controller.signal;
+
+    const errors = [];
+    function onError(error) {
+      errors.push(error);
+    }
+    const postpones = [];
+    function onPostpone(reason) {
+      postpones.push(reason);
+    }
+    let pendingPrerender;
+    await act(() => {
+      pendingPrerender = ReactDOMFizzStatic.prerenderToNodeStream(<App />, {
+        signal,
+        onError,
+        onPostpone,
+      });
+    });
+    controller.abort(postponeInstance);
+
+    const prerendered = await pendingPrerender;
+
+    expect(prerendered.postponed).toBe(null);
+    expect(errors).toEqual([]);
+    expect(postpones).toEqual(['manufactured', 'manufactured']);
+
+    await act(() => {
+      prerendered.prelude.pipe(writable);
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <p>
+          <span>Loading again...</span>
+        </p>
+        <p>
+          <span>Loading again too...</span>
+        </p>
+      </div>,
+    );
+  });
+
+  // @gate enablePostpone
+  it('does not call onError when you abort with a postpone instance during resume', async () => {
+    let prerendering = true;
+    const promise = new Promise(r => {});
+
+    function Wait() {
+      return React.use(promise);
+    }
+    function Postpone() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return (
+        <span>
+          <Suspense fallback="Loading again...">
+            <Wait />
+          </Suspense>
+        </span>
+      );
+    }
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <p>
+              <Postpone />
+            </p>
+            <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);
+    });
+
+    expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
+
+    let postponeInstance;
+    try {
+      React.unstable_postpone('manufactured');
+    } catch (p) {
+      postponeInstance = p;
+    }
+
+    const errors = [];
+    function onError(error) {
+      errors.push(error);
+    }
+    const postpones = [];
+    function onPostpone(reason) {
+      postpones.push(reason);
+    }
+
+    prerendering = false;
+
+    const resumed = await ReactDOMFizzServer.resumeToPipeableStream(
+      <App />,
+      JSON.parse(JSON.stringify(prerendered.postponed)),
+      {
+        onError,
+        onPostpone,
+      },
+    );
+
+    await act(() => {
+      resumed.pipe(writable);
+    });
+    await act(() => {
+      resumed.abort(postponeInstance);
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <p>
+          <span>Loading again...</span>
+        </p>
+        <p>
+          <span>Loading again...</span>
+        </p>
+      </div>,
+    );
+
+    expect(errors).toEqual([]);
+    expect(postpones).toEqual(['manufactured', 'manufactured']);
+  });
+
+  // @gate enablePostpone
+  it('does not call onError when you abort with a postpone instance during a render', async () => {
+    const promise = new Promise(r => {});
+
+    function Wait() {
+      return React.use(promise);
+    }
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <p>
+              <span>
+                <Suspense fallback="Loading again...">
+                  <Wait />
+                </Suspense>
+              </span>
+            </p>
+            <p>
+              <span>
+                <Suspense fallback="Loading again...">
+                  <Wait />
+                </Suspense>
+              </span>
+            </p>
+          </Suspense>
+        </div>
+      );
+    }
+
+    const errors = [];
+    function onError(error) {
+      errors.push(error);
+    }
+    const postpones = [];
+    function onPostpone(reason) {
+      postpones.push(reason);
+    }
+    const result = await renderToPipeableStream(<App />, {onError, onPostpone});
+    await act(() => {
+      result.pipe(writable);
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <p>
+          <span>Loading again...</span>
+        </p>
+        <p>
+          <span>Loading again...</span>
+        </p>
+      </div>,
+    );
+
+    let postponeInstance;
+    try {
+      React.unstable_postpone('manufactured');
+    } catch (p) {
+      postponeInstance = p;
+    }
+    await act(() => {
+      result.abort(postponeInstance);
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <p>
+          <span>Loading again...</span>
+        </p>
+        <p>
+          <span>Loading again...</span>
+        </p>
+      </div>,
+    );
+
+    expect(errors).toEqual([]);
+    expect(postpones).toEqual(['manufactured', 'manufactured']);
+  });
+
+  // @gate enablePostpone
+  it('fatally errors if you abort with a postpone in the shell during resume', async () => {
+    let prerendering = true;
+    const promise = new Promise(r => {});
+
+    function Wait() {
+      return React.use(promise);
+    }
+    function Postpone() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return (
+        <span>
+          <Suspense fallback="Loading again...">
+            <Wait />
+          </Suspense>
+        </span>
+      );
+    }
+
+    function PostponeInShell() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return <span>in shell</span>;
+    }
+
+    function App() {
+      return (
+        <div>
+          <PostponeInShell />
+          <Suspense fallback="Loading...">
+            <p>
+              <Postpone />
+            </p>
+            <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);
+    });
+
+    expect(getVisibleChildren(container)).toEqual(undefined);
+
+    let postponeInstance;
+    try {
+      React.unstable_postpone('manufactured');
+    } catch (p) {
+      postponeInstance = p;
+    }
+
+    const errors = [];
+    function onError(error) {
+      errors.push(error);
+    }
+    const shellErrors = [];
+    function onShellError(error) {
+      shellErrors.push(error);
+    }
+    const postpones = [];
+    function onPostpone(reason) {
+      postpones.push(reason);
+    }
+
+    prerendering = false;
+
+    const resumed = await ReactDOMFizzServer.resumeToPipeableStream(
+      <App />,
+      JSON.parse(JSON.stringify(prerendered.postponed)),
+      {
+        onError,
+        onShellError,
+        onPostpone,
+      },
+    );
+    await act(() => {
+      resumed.abort(postponeInstance);
+    });
+    expect(errors).toEqual([
+      new Error(
+        'The render was aborted with postpone when the shell is incomplete. Reason: manufactured',
+      ),
+    ]);
+    expect(shellErrors).toEqual([
+      new Error(
+        'The render was aborted with postpone when the shell is incomplete. Reason: manufactured',
+      ),
+    ]);
+    expect(postpones).toEqual([]);
+  });
+
+  // @gate enablePostpone
+  it('fatally errors if you abort with a postpone in the shell during render', async () => {
+    const promise = new Promise(r => {});
+
+    function Wait() {
+      return React.use(promise);
+    }
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <p>
+              <span>
+                <Suspense fallback="Loading again...">
+                  <Wait />
+                </Suspense>
+              </span>
+            </p>
+            <p>
+              <span>
+                <Suspense fallback="Loading again...">
+                  <Wait />
+                </Suspense>
+              </span>
+            </p>
+          </Suspense>
+        </div>
+      );
+    }
+
+    const errors = [];
+    function onError(error) {
+      errors.push(error);
+    }
+    const shellErrors = [];
+    function onShellError(error) {
+      shellErrors.push(error);
+    }
+    const postpones = [];
+    function onPostpone(reason) {
+      postpones.push(reason);
+    }
+    const result = await renderToPipeableStream(<App />, {
+      onError,
+      onShellError,
+      onPostpone,
+    });
+
+    let postponeInstance;
+    try {
+      React.unstable_postpone('manufactured');
+    } catch (p) {
+      postponeInstance = p;
+    }
+    await act(() => {
+      result.abort(postponeInstance);
+    });
+
+    expect(getVisibleChildren(container)).toEqual(undefined);
+
+    expect(errors).toEqual([
+      new Error(
+        'The render was aborted with postpone when the shell is incomplete. Reason: manufactured',
+      ),
+    ]);
+    expect(shellErrors).toEqual([
+      new Error(
+        'The render was aborted with postpone when the shell is incomplete. Reason: manufactured',
+      ),
+    ]);
+    expect(postpones).toEqual([]);
+  });
 });
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index 6b82bfe3a45b7..152b3acc24593 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -3129,8 +3129,23 @@ function abortTask(task: Task, request: Request, error: mixed): void {
       if (replay === null) {
         // We didn't complete the root so we have nothing to show. We can close
         // the request;
-        logRecoverableError(request, error, errorInfo);
-        fatalError(request, error);
+        if (
+          enablePostpone &&
+          typeof error === 'object' &&
+          error !== null &&
+          error.$$typeof === REACT_POSTPONE_TYPE
+        ) {
+          const postponeInstance: Postpone = (error: any);
+          const fatal = new Error(
+            'The render was aborted with postpone when the shell is incomplete. Reason: ' +
+              postponeInstance.message,
+          );
+          logRecoverableError(request, fatal, errorInfo);
+          fatalError(request, fatal);
+        } else {
+          logRecoverableError(request, error, errorInfo);
+          fatalError(request, error);
+        }
         return;
       } else {
         // If the shell aborts during a replay, that's not a fatal error. Instead
@@ -3138,7 +3153,20 @@ function abortTask(task: Task, request: Request, error: mixed): void {
         // the ReplaySet.
         replay.pendingTasks--;
         if (replay.pendingTasks === 0 && replay.nodes.length > 0) {
-          const errorDigest = logRecoverableError(request, error, errorInfo);
+          let errorDigest;
+          if (
+            enablePostpone &&
+            typeof error === 'object' &&
+            error !== null &&
+            error.$$typeof === REACT_POSTPONE_TYPE
+          ) {
+            const postponeInstance: Postpone = (error: any);
+            logPostpone(request, postponeInstance.message, errorInfo);
+            // TODO: Figure out a better signal than a magic digest value.
+            errorDigest = 'POSTPONE';
+          } else {
+            errorDigest = logRecoverableError(request, error, errorInfo);
+          }
           abortRemainingReplayNodes(
             request,
             null,
@@ -3162,7 +3190,20 @@ function abortTask(task: Task, request: Request, error: mixed): void {
       // We construct an errorInfo from the boundary's componentStack so the error in dev will indicate which
       // boundary the message is referring to
       const errorInfo = getThrownInfo(request, task.componentStack);
-      const errorDigest = logRecoverableError(request, error, errorInfo);
+      let errorDigest;
+      if (
+        enablePostpone &&
+        typeof error === 'object' &&
+        error !== null &&
+        error.$$typeof === REACT_POSTPONE_TYPE
+      ) {
+        const postponeInstance: Postpone = (error: any);
+        logPostpone(request, postponeInstance.message, errorInfo);
+        // TODO: Figure out a better signal than a magic digest value.
+        errorDigest = 'POSTPONE';
+      } else {
+        errorDigest = logRecoverableError(request, error, errorInfo);
+      }
       let errorMessage = error;
       if (__DEV__) {
         const errorPrefix =
diff --git a/scripts/error-codes/codes.json b/scripts/error-codes/codes.json
index 74ae5ea927dc6..ad5d3ebaf3f5e 100644
--- a/scripts/error-codes/codes.json
+++ b/scripts/error-codes/codes.json
@@ -485,5 +485,6 @@
   "497": "Only objects or functions can be passed to taintObjectReference.",
   "498": "Only plain objects, and a few built-ins, can be passed to Client Components from Server Components. Classes or null prototypes are not supported.",
   "499": "Only plain objects, and a few built-ins, can be passed to Server Actions. Classes or null prototypes are not supported.",
-  "500": "React expected a headers state to exist when emitEarlyPreloads was called but did not find it. This suggests emitEarlyPreloads was called more than once per request. This is a bug in React."
+  "500": "React expected a headers state to exist when emitEarlyPreloads was called but did not find it. This suggests emitEarlyPreloads was called more than once per request. This is a bug in React.",
+  "501": "The render was aborted with postpone when the shell is incomplete. Reason: %s"
 }