From de184c399d2bcd53a1baa079c886056ad926afde Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Fri, 16 Aug 2024 15:07:41 -0700
Subject: [PATCH] [Fizz] handle throwing after abort during render

It is possible to throw after aborting during a render and we were not properly tracking this. We use an AbortSigil to mark whether a rendering task needs to abort but the throw interrupts that and we end up handling an error on the error pathway instead.

This change reworks the abort-while-rendering support to be robust to throws after calling abort
---
 .../src/__tests__/ReactDOMFizzServer-test.js  | 37 +++++++++++++++++++
 packages/react-server/src/ReactFizzServer.js  | 22 ++++++-----
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
index e9abb25eb1768..b6985c403f34c 100644
--- a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
+++ b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
@@ -8377,6 +8377,43 @@ describe('ReactDOMFizzServer', () => {
     );
   });
 
+  it('can support throwing after aborting during a render', async () => {
+    function App() {
+      return (
+        <div>
+          <Suspense fallback={<p>loading...</p>}>
+            <ComponentThatAborts />
+          </Suspense>
+        </div>
+      );
+    }
+
+    function ComponentThatAborts() {
+      abortRef.current('boom');
+      throw new Error('bam');
+    }
+
+    const abortRef = {current: null};
+    let finished = false;
+    await act(() => {
+      const {pipe, abort} = renderToPipeableStream(<App />);
+      abortRef.current = abort;
+      writable.on('finish', () => {
+        finished = true;
+      });
+      pipe(writable);
+    });
+
+    assertConsoleErrorDev(['boom']);
+
+    expect(finished).toBe(true);
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <p>loading...</p>
+      </div>,
+    );
+  });
+
   it('should warn for using generators as children props', async () => {
     function* getChildren() {
       yield <h1 key="1">Hello</h1>;
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index a6c77f8bafb7f..986e8673a5a2a 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -652,8 +652,6 @@ export function resumeRequest(
   return request;
 }
 
-const AbortSigil = {};
-
 let currentRequest: null | Request = null;
 
 export function resolveRequest(): null | Request {
@@ -1173,7 +1171,7 @@ function renderSuspenseBoundary(
       );
       boundarySegment.status = COMPLETED;
     } catch (thrownValue: mixed) {
-      if (thrownValue === AbortSigil) {
+      if (request.status === ABORTING) {
         boundarySegment.status = ABORTED;
       } else {
         boundarySegment.status = ERRORED;
@@ -1246,7 +1244,7 @@ function renderSuspenseBoundary(
     } catch (thrownValue: mixed) {
       newBoundary.status = CLIENT_RENDERED;
       let error: mixed;
-      if (thrownValue === AbortSigil) {
+      if (request.status === ABORTING) {
         contentRootSegment.status = ABORTED;
         error = request.fatalError;
       } else {
@@ -1601,7 +1599,8 @@ function finishClassComponent(
     nextChildren = instance.render();
   }
   if (request.status === ABORTING) {
-    throw AbortSigil;
+    // eslint-disable-next-line no-throw-literal
+    throw null;
   }
 
   if (__DEV__) {
@@ -1757,7 +1756,8 @@ function renderFunctionComponent(
     legacyContext,
   );
   if (request.status === ABORTING) {
-    throw AbortSigil;
+    // eslint-disable-next-line no-throw-literal
+    throw null;
   }
 
   const hasId = checkDidRenderIdHook();
@@ -2076,7 +2076,8 @@ function renderLazyComponent(
     Component = init(payload);
   }
   if (request.status === ABORTING) {
-    throw AbortSigil;
+    // eslint-disable-next-line no-throw-literal
+    throw null;
   }
   const resolvedProps = resolveDefaultPropsOnNonClassComponent(
     Component,
@@ -2655,7 +2656,8 @@ function retryNode(request: Request, task: Task): void {
           resolvedNode = init(payload);
         }
         if (request.status === ABORTING) {
-          throw AbortSigil;
+          // eslint-disable-next-line no-throw-literal
+          throw null;
         }
         // Now we render the resolved node
         renderNodeDestructive(request, task, resolvedNode, childIndex);
@@ -4127,7 +4129,7 @@ function retryRenderTask(
           // (unstable) API for suspending. This implementation detail can change
           // later, once we deprecate the old API in favor of `use`.
           getSuspendedThenable()
-        : thrownValue === AbortSigil
+        : request.status === ABORTING
           ? request.fatalError
           : thrownValue;
 
@@ -4250,7 +4252,7 @@ function retryReplayTask(request: Request, task: ReplayTask): void {
     erroredReplay(
       request,
       task.blockedBoundary,
-      x === AbortSigil ? request.fatalError : x,
+      request.status === ABORTING ? request.fatalError : x,
       errorInfo,
       task.replay.nodes,
       task.replay.slots,