[Fizz] Abort tasks that suspend after aborting during render

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

@@ -7138,7 +7138,7 @@ describe('ReactDOMFizzServer', () => {
     expect(errors).toEqual(['abort reason', 'abort reason']);
   });
 
-  it('currently does not report a root task that suspends directly after aborting during render', async () => {
+  it('reports a root task that suspends directly after aborting during render', async () => {
     const promise = new Promise(() => {});
     const abortRef = {current: null};
     function ComponentThatAbortsAndSuspends() {
@@ -7161,9 +7161,7 @@ describe('ReactDOMFizzServer', () => {
       abortRef.current = abort;
     });
 
-    // The task suspends before renderFunctionComponent gets to throw the
-    // abort reason, and retryRenderTask currently suspends it again.
-    expect(errors).toEqual([]);
+    expect(errors).toEqual(['abort reason']);
   });
 
   it('can abort during render in a lazy initializer for a component', async () => {

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

@@ -379,6 +379,29 @@ describe('ReactDOMFizzStaticBrowser', () => {
     expect(content).toBe('');
   });
 
+  it('reports the abort reason if a task suspends after aborting a prerender', async () => {
+    const promise = new Promise(() => {});
+    const errors = [];
+    const controller = new AbortController();
+    function App() {
+      controller.abort(new Error('abort reason'));
+      React.use(promise);
+      return null;
+    }
+
+    const result = await serverAct(() =>
+      ReactDOMFizzStatic.prerender(<App />, {
+        signal: controller.signal,
+        onError(error) {
+          errors.push(error.message);
+        },
+      }),
+    );
+
+    expect(errors).toEqual(['abort reason']);
+    expect(await readContent(result.prelude)).toBe('');
+  });
+
   it('should resolve an empty prelude if passing an already aborted signal', async () => {
     const errors = [];
     const controller = new AbortController();
@@ -724,6 +747,66 @@ describe('ReactDOMFizzStaticBrowser', () => {
     expect(errors).toEqual(['resume abort']);
   });
 
+  it('can abort and suspend while replaying a prerendered tree', async () => {
+    const promise = new Promise(() => {});
+    let prerendering = true;
+    const resumeController = new AbortController();
+
+    function AbortDuringReplay({children}) {
+      if (!prerendering) {
+        resumeController.abort('resume abort');
+        React.use(promise);
+      }
+      return children;
+    }
+
+    function Wait() {
+      return React.use(promise);
+    }
+
+    function App() {
+      return (
+        <div>
+          <AbortDuringReplay>
+            <Suspense fallback="Loading...">
+              <Wait />
+            </Suspense>
+          </AbortDuringReplay>
+        </div>
+      );
+    }
+
+    const controller = new AbortController();
+    let pendingResult;
+    await serverAct(() => {
+      pendingResult = ReactDOMFizzStatic.prerender(<App />, {
+        signal: controller.signal,
+        onError() {},
+      });
+    });
+    await serverAct(() => {
+      controller.abort('prerender abort');
+    });
+    const prerendered = await pendingResult;
+
+    prerendering = false;
+    const errors = [];
+    await serverAct(() =>
+      ReactDOMFizzServer.resume(
+        <App />,
+        JSON.parse(JSON.stringify(prerendered.postponed)),
+        {
+          signal: resumeController.signal,
+          onError(error) {
+            errors.push(error);
+          },
+        },
+      ),
+    );
+
+    expect(errors).toEqual(['resume abort']);
+  });
+
   it('can abort while rendering a resumed segment', async () => {
     const promise = new Promise(() => {});
     let prerendering = true;

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

@@ -14,6 +14,35 @@ let React;
 let ReactDOMFizzStatic;
 let Suspense;
 
+function normalizeCodeLocInfo(str) {
+  return (
+    str &&
+    str.replace(/^ +(?:at|in) ([\S]+)[^\n]*/gm, function (m, name) {
+      const dot = name.lastIndexOf('.');
+      if (dot !== -1) {
+        name = name.slice(dot + 1);
+      }
+      return '    in ' + name + (/\d/.test(m) ? ' (at **)' : '');
+    })
+  );
+}
+
+function ignoreListStack(str) {
+  if (!str) {
+    return str;
+  }
+
+  let ignoreListedStack = '';
+  const lines = str.split('\n');
+  // eslint-disable-next-line no-for-of-loops/no-for-of-loops
+  for (const line of lines) {
+    if (line.indexOf(__filename) !== -1) {
+      ignoreListedStack += '\n' + line;
+    }
+  }
+  return ignoreListedStack;
+}
+
 describe('ReactDOMFizzStaticNode', () => {
   beforeEach(() => {
     jest.resetModules();
@@ -290,6 +319,30 @@ describe('ReactDOMFizzStaticNode', () => {
     expect(content).toBe('');
   });
 
+  it('reports the abort reason if a task suspends after aborting a prerender', async () => {
+    const promise = new Promise(() => {});
+    const errors = [];
+    const controller = new AbortController();
+    function App() {
+      controller.abort(new Error('abort reason'));
+      React.use(promise);
+      return null;
+    }
+
+    const resultPromise = ReactDOMFizzStatic.prerenderToNodeStream(<App />, {
+      signal: controller.signal,
+      onError(error) {
+        errors.push(error.message);
+      },
+    });
+
+    await jest.runAllTimers();
+    const result = await resultPromise;
+
+    expect(errors).toEqual(['abort reason']);
+    expect(await readContent(result.prelude)).toBe('');
+  });
+
   it('should resolve with an empty prelude if passing an already aborted signal', async () => {
     const errors = [];
     const controller = new AbortController();
@@ -457,4 +510,61 @@ describe('ReactDOMFizzStaticNode', () => {
 
     expect(errors).toEqual(['abort reason', 'abort reason']);
   });
+
+  describe('with real timers', () => {
+    beforeEach(() => {
+      jest.useRealTimers();
+    });
+
+    afterEach(() => {
+      jest.useFakeTimers();
+    });
+
+    it('includes the suspended call site when aborting in the same rendering task', async () => {
+      const promise = new Promise(() => {});
+      const controller = new AbortController();
+      let caughtError;
+      let componentStack;
+      let ownerStack;
+
+      function AbortAndSuspend() {
+        controller.abort(new Error('abort reason'));
+        React.use(promise);
+        return null;
+      }
+
+      function App() {
+        return <AbortAndSuspend />;
+      }
+
+      const {prelude} = await ReactDOMFizzStatic.prerenderToNodeStream(
+        <App />,
+        {
+          signal: controller.signal,
+          onError(error, errorInfo) {
+            caughtError = error;
+            componentStack = errorInfo.componentStack;
+            ownerStack = __DEV__ ? React.captureOwnerStack() : null;
+          },
+        },
+      );
+
+      expect(caughtError).toEqual(
+        expect.objectContaining({message: 'abort reason'}),
+      );
+      expect(await readContent(prelude)).toBe('');
+      if (__DEV__) {
+        expect(normalizeCodeLocInfo(componentStack)).toBe(
+          '\n    in AbortAndSuspend (at **)\n    in App',
+        );
+        expect(normalizeCodeLocInfo(ignoreListStack(ownerStack))).toBe(
+          (gate(flags => flags.enableAsyncDebugInfo)
+            ? '\n    in AbortAndSuspend (at **)'
+            : '') + '\n    in App (at **)',
+        );
+      } else {
+        expect(ownerStack).toBeNull();
+      }
+    });
+  });
 });

packages/react-server/src/ReactFizzServer.js

@@ -4829,6 +4829,23 @@ function abortTaskDEV(task: Task, request: Request): void {
   }
 }
 
+function abortUnwoundTask(task: Task, request: Request): void {
+  // This task was rendering when abort began, so the synchronous abort sweep
+  // left it alone. It has now unwound from user code and can be completed
+  // through the normal abort path.
+  if (__DEV__) {
+    abortTaskDEV(task, request);
+  } else {
+    abortTask(task, request);
+  }
+  task.abortSet.delete(task);
+  if (__DEV__) {
+    finishAbortedTaskDEV(task, request, request.fatalError);
+  } else {
+    finishAbortedTask(task, request, request.fatalError);
+  }
+}
+
 function safelyEmitEarlyPreloads(
   request: Request,
   shellComplete: boolean,
@@ -5169,25 +5186,22 @@ function retryRenderTask(
           // (unstable) API for suspending. This implementation detail can change
           // later, once we deprecate the old API in favor of `use`.
           getSuspendedThenable()
-        : request.aborted
-          ? request.fatalError
-          : thrownValue;
-
-    if (request.aborted && request.trackedPostpones !== null) {
-      // We are aborting a prerender and need to halt this task.
-      const trackedPostpones = request.trackedPostpones;
-      const thrownInfo = getThrownInfo(task.componentStack);
-      task.abortSet.delete(task);
-
-      logRecoverableError(
-        request,
-        x,
-        thrownInfo,
-        __DEV__ ? task.debugTask : null,
-      );
+        : thrownValue;
 
-      trackPostpone(request, trackedPostpones, task, segment);
-      finishedTask(request, task.blockedBoundary, task.row, segment);
+    if (request.aborted) {
+      if (thrownValue === SuspenseException) {
+        // This task was rendering when abort() was called, so it never took
+        // the normal suspension path below that stores the thenable state.
+        // Preserve it before finishing the abort so DEV can replay the task
+        // and include this suspended use() call site in the owner stack.
+        task.thenableState = getThenableStateAfterSuspending();
+      }
+      // The task has unwound from user code, so it must no longer appear to
+      // be the currently rendering task while we synchronously finish it.
+      // Restore the parent instead of clearing this field because finishing
+      // can reenter Fizz and abort an outer render that is still on the stack.
+      request.currentTask = prevTask;
+      abortUnwoundTask(task, request);
       return;
     }
 
@@ -5280,6 +5294,23 @@ function retryReplayTask(request: Request, task: ReplayTask): void {
           getSuspendedThenable()
         : thrownValue;
 
+    if (request.aborted) {
+      if (thrownValue === SuspenseException) {
+        // This task was rendering when abort() was called, so it never took
+        // the normal suspension path below that stores the thenable state.
+        // Preserve it before finishing the abort so DEV can replay the task
+        // and include this suspended use() call site in the owner stack.
+        task.thenableState = getThenableStateAfterSuspending();
+      }
+      // The task has unwound from user code, so it must no longer appear to
+      // be the currently rendering task while we synchronously finish it.
+      // Restore the parent instead of clearing this field because finishing
+      // can reenter Fizz and abort an outer render that is still on the stack.
+      request.currentTask = prevTask;
+      abortUnwoundTask(task, request);
+      return;
+    }
+
     if (typeof x === 'object' && x !== null) {
       // $FlowFixMe[method-unbinding]
       if (typeof x.then === 'function') {