[Fizz] Allow Pending Work to Specialize Abort Reasons

packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js

@@ -3603,6 +3603,88 @@ describe('ReactDOMFizzServer', () => {
     expect(errors).toEqual(['abort reason', 'abort reason', 'abort reason']);
   });
 
+  it('uses a rejection reason from a lazy component before the abort finishes', async () => {
+    let reject;
+    const Lazy = React.lazy(
+      () =>
+        new Promise((resolve, rejectPromise) => {
+          reject = rejectPromise;
+        }),
+    );
+    const haltedPromise = new Promise(() => {});
+    function HaltedWait() {
+      use(haltedPromise);
+      return null;
+    }
+
+    const errors = [];
+    let abort;
+    await act(() => {
+      const controls = renderToPipeableStream(
+        <>
+          <Suspense fallback="Loading lazy">
+            <Lazy />
+          </Suspense>
+          <Suspense fallback="Loading halted">
+            <HaltedWait />
+          </Suspense>
+        </>,
+        {
+          onError(error) {
+            errors.push(error.message);
+          },
+        },
+      );
+      abort = controls.abort;
+      controls.pipe(writable);
+    });
+
+    await act(() => {
+      abort(new Error('abort reason'));
+      reject(new Error('rejected during abort'));
+    });
+
+    expect(errors).toEqual(['rejected during abort', 'abort reason']);
+  });
+
+  it('does not report a rejection reason after abort has finished', async () => {
+    let reject;
+    const promise = new Promise((resolve, rejectPromise) => {
+      reject = rejectPromise;
+    });
+    function Wait() {
+      use(promise);
+      return null;
+    }
+
+    const errors = [];
+    let abort;
+    await act(() => {
+      const controls = renderToPipeableStream(
+        <Suspense fallback="Loading">
+          <Wait />
+        </Suspense>,
+        {
+          onError(error) {
+            errors.push(error.message);
+          },
+        },
+      );
+      abort = controls.abort;
+      controls.pipe(writable);
+    });
+
+    await act(() => {
+      abort(new Error('abort reason'));
+    });
+
+    await act(() => {
+      reject(new Error('rejected after abort'));
+    });
+
+    expect(errors).toEqual(['abort reason']);
+  });
+
   it('warns in dev if you access digest from errorInfo in onRecoverableError', async () => {
     await act(() => {
       const {pipe} = renderToPipeableStream(

packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js

@@ -527,7 +527,7 @@ describe('ReactDOMFizzStaticBrowser', () => {
     expect(errors).toEqual(['uh oh', 'uh oh']);
   });
 
-  it('currently uses the abort reason when an abort listener synchronously rejects pending work', async () => {
+  it('uses a rejection reason when an abort listener rejects pending work before the abort finishes', async () => {
     let reject;
     const rejectedPromise = new Promise((resolve, rejectPromise) => {
       reject = rejectPromise;
@@ -572,7 +572,7 @@ describe('ReactDOMFizzStaticBrowser', () => {
     });
     await resultPromise;
 
-    expect(errors).toEqual(['abort reason', 'abort reason']);
+    expect(errors).toEqual(['rejected during abort', 'abort reason']);
   });
 
   it('logs an error if onHeaders throws but continues the prerender', async () => {

packages/react-dom/src/__tests__/ReactDOMFizzStaticNode-test.js

@@ -464,7 +464,7 @@ describe('ReactDOMFizzStaticNode', () => {
     expect(errors).toEqual(['uh oh', 'uh oh']);
   });
 
-  it('currently uses the abort reason when an abort listener synchronously rejects pending work', async () => {
+  it('uses a rejection reason when an abort listener rejects pending work before the abort finishes', async () => {
     let reject;
     const rejectedPromise = new Promise((resolve, rejectPromise) => {
       reject = rejectPromise;
@@ -505,10 +505,11 @@ describe('ReactDOMFizzStaticNode', () => {
     });
     controller.abort(new Error('abort reason'));
 
+    await Promise.resolve();
     await jest.runAllTimers();
     await resultPromise;
 
-    expect(errors).toEqual(['abort reason', 'abort reason']);
+    expect(errors).toEqual(['rejected during abort', 'abort reason']);
   });
 
   describe('with real timers', () => {

packages/react-server/src/ReactFizzServer.js

@@ -276,11 +276,16 @@ type SuspenseBoundary = {
   errorComponentStack?: null | string, // the error component stack if it errors
 };
 
+type Ping = {
+  resolve: () => void,
+  reject: (error: mixed) => void,
+};
+
 type RenderTask = {
   replay: null,
   node: ReactNodeList,
   childIndex: number,
-  ping: () => void,
+  ping: Ping,
   blockedBoundary: Root | SuspenseBoundary,
   blockedSegment: Segment, // the segment we'll write to
   blockedPreamble: null | PreambleState,
@@ -311,7 +316,7 @@ type ReplayTask = {
   replay: ReplaySet,
   node: ReactNodeList,
   childIndex: number,
-  ping: () => void,
+  ping: Ping,
   blockedBoundary: Root | SuspenseBoundary,
   blockedSegment: null, // we don't write to anything when we replay
   blockedPreamble: null,
@@ -810,6 +815,27 @@ function pingTask(request: Request, task: Task): void {
   }
 }
 
+function pingRejectedTask(request: Request, task: Task, error: mixed): void {
+  if (!request.aborted) {
+    // Replaying the task is what gives ordinary render errors their complete
+    // component stack.
+    pingTask(request, task);
+    return;
+  }
+  if (!task.abortSet.delete(task)) {
+    // finishAbort already completed this task with the request's abort reason.
+    return;
+  }
+  // abortTask synchronously claimed this task before abort listeners could
+  // reject its wakeable. Finish it with the more specific reason before the
+  // scheduled final abort uses the reason for the whole request.
+  if (__DEV__) {
+    finishAbortedTaskDEV(task, request, error);
+  } else {
+    finishAbortedTask(task, request, error);
+  }
+}
+
 function createSuspenseBoundary(
   request: Request,
   row: null | SuspenseListRow,
@@ -889,7 +915,10 @@ function createRenderTask(
     replay: null,
     node,
     childIndex,
-    ping: () => pingTask(request, task),
+    ping: {
+      resolve: () => pingTask(request, task),
+      reject: error => pingRejectedTask(request, task, error),
+    },
     blockedBoundary,
     blockedSegment,
     blockedPreamble,
@@ -945,7 +974,10 @@ function createReplayTask(
     replay,
     node,
     childIndex,
-    ping: () => pingTask(request, task),
+    ping: {
+      resolve: () => pingTask(request, task),
+      reject: error => pingRejectedTask(request, task, error),
+    },
     blockedBoundary,
     blockedSegment: null,
     blockedPreamble: null,
@@ -4204,7 +4236,7 @@ function renderNode(
             thenableState,
           );
           const ping = newTask.ping;
-          wakeable.then(ping, ping);
+          wakeable.then(ping.resolve, ping.reject);
 
           // Restore the context. We assume that this will be restored by the inner
           // functions in case nothing throws so we don't use "finally" here.
@@ -4305,7 +4337,7 @@ function renderNode(
             thenableState,
           );
           const ping = newTask.ping;
-          wakeable.then(ping, ping);
+          wakeable.then(ping.resolve, ping.reject);
 
           // Restore the context. We assume that this will be restored by the inner
           // functions in case nothing throws so we don't use "finally" here.
@@ -5216,7 +5248,7 @@ function retryRenderTask(
             : null;
         const ping = task.ping;
         // We've asserted that x is a thenable above
-        (x: any).then(ping, ping);
+        (x: any).then(ping.resolve, ping.reject);
         return;
       }
     }
@@ -5316,7 +5348,7 @@ function retryReplayTask(request: Request, task: ReplayTask): void {
       if (typeof x.then === 'function') {
         // Something suspended again, let's pick it back up later.
         const ping = task.ping;
-        x.then(ping, ping);
+        x.then(ping.resolve, ping.reject);
         task.thenableState =
           thrownValue === SuspenseException
             ? getThenableStateAfterSuspending()