[Fizz] Finish abort in a scheduled task

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

@@ -7063,7 +7063,7 @@ describe('ReactDOMFizzServer', () => {
     );
   });
 
-  it('currently does not report an in-flight root task after another root task fatals while aborting', async () => {
+  it('reports an in-flight root task after another root task fatals while aborting', async () => {
     const promise = new Promise(() => {});
     function SuspendedRoot() {
       use(promise);
@@ -7097,10 +7097,10 @@ describe('ReactDOMFizzServer', () => {
       abortRef.current = abort;
     });
 
-    expect(errors).toEqual(['abort reason']);
+    expect(errors).toEqual(['abort reason', 'abort reason']);
   });
 
-  it('currently does not report a root task that suspends after aborting during render', async () => {
+  it('reports a root task before rendering a suspended child returned after aborting', async () => {
     const promise = new Promise(() => {});
     function SuspendedRoot() {
       use(promise);
@@ -7135,9 +7135,35 @@ describe('ReactDOMFizzServer', () => {
       abortRef.current = abort;
     });
 
-    // TODO: Once abort completion is async, this still-suspended task should
-    // observe ABORTING and report the abort reason as well.
-    expect(errors).toEqual(['abort reason']);
+    expect(errors).toEqual(['abort reason', 'abort reason']);
+  });
+
+  it('currently does not report a root task that suspends directly after aborting during render', async () => {
+    const promise = new Promise(() => {});
+    const abortRef = {current: null};
+    function ComponentThatAbortsAndSuspends() {
+      abortRef.current(new Error('abort reason'));
+      use(promise);
+      return null;
+    }
+
+    const errors = [];
+    await act(() => {
+      const {abort} = renderToPipeableStream(
+        <ComponentThatAbortsAndSuspends />,
+        {
+          onError(error) {
+            errors.push(error.message);
+          },
+          onShellError() {},
+        },
+      );
+      abortRef.current = abort;
+    });
+
+    // The task suspends before renderFunctionComponent gets to throw the
+    // abort reason, and retryRenderTask currently suspends it again.
+    expect(errors).toEqual([]);
   });
 
   it('can abort during render in a lazy initializer for a component', async () => {
@@ -8444,11 +8470,11 @@ describe('ReactDOMFizzServer', () => {
     }
 
     expect(thrownError).toBe('boom');
-    // TODO there should actually be three errors. One for the pending Suspense, one for the fallback task, and one for the task
-    // that does the abort itself. At the moment abort will flush queues and if there is no pending tasks will close the request before
-    // the task which initiated the abort can even be processed. This is a bug but not one that I am fixing with the current change
-    // so I am asserting the current behavior
     expect(errors).toEqual([
+      {
+        error: 'boom',
+        componentStack: componentStack(['Abort', 'body', 'html', 'App']),
+      },
       {
         error: 'boom',
         componentStack: componentStack([
@@ -8467,9 +8493,6 @@ describe('ReactDOMFizzServer', () => {
           'html',
           'App',
         ]),
-        // }, {
-        //   error: 'boom',
-        //   componentStack: componentStack(['Abort', 'body', 'html', 'App'])
       },
     ]);
 

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

@@ -221,7 +221,9 @@ describe('ReactDOMFizzServerBrowser', () => {
       ),
     );
 
-    controller.abort();
+    await serverAct(() => {
+      controller.abort();
+    });
 
     const result = await readResult(stream);
     expect(result).toContain('Loading');
@@ -247,7 +249,9 @@ describe('ReactDOMFizzServerBrowser', () => {
     );
 
     const theReason = new Error('aborted for reasons');
-    controller.abort(theReason);
+    await serverAct(() => {
+      controller.abort(theReason);
+    });
 
     let caughtError = null;
     try {
@@ -364,7 +368,7 @@ describe('ReactDOMFizzServerBrowser', () => {
 
     const reader = stream.getReader();
     await reader.read();
-    await reader.cancel();
+    await serverAct(() => reader.cancel());
 
     expect(errors).toEqual([
       'The render was aborted by the server without a reason.',
@@ -463,7 +467,9 @@ describe('ReactDOMFizzServerBrowser', () => {
       }),
     );
 
-    controller.abort('foobar');
+    await serverAct(() => {
+      controller.abort('foobar');
+    });
 
     expect(errors).toEqual(['foobar', 'foobar']);
   });
@@ -502,7 +508,9 @@ describe('ReactDOMFizzServerBrowser', () => {
       }),
     );
 
-    controller.abort(new Error('uh oh'));
+    await serverAct(() => {
+      controller.abort(new Error('uh oh'));
+    });
 
     expect(errors).toEqual(['uh oh', 'uh oh']);
   });

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

@@ -382,6 +382,7 @@ describe('ReactDOMFizzServerNode', () => {
     expect(isCompleteCalls).toBe(0);
 
     abort(new Error('uh oh'));
+    await jest.runAllTimers();
 
     await completed;
 
@@ -468,6 +469,7 @@ describe('ReactDOMFizzServerNode', () => {
 
     const reason = new Error('abort reason');
     abort(reason);
+    await jest.runAllTimers();
 
     expect(shellErrors).toEqual([reason]);
     expect(errors).toEqual(['abort reason', 'abort reason', 'abort reason']);
@@ -504,6 +506,7 @@ describe('ReactDOMFizzServerNode', () => {
     expect(isCompleteCalls).toBe(0);
 
     abort();
+    await jest.runAllTimers();
 
     await completed;
 
@@ -746,6 +749,7 @@ describe('ReactDOMFizzServerNode', () => {
     resolve();
 
     await completed;
+    await jest.runAllTimers();
 
     expect(errors).toEqual([
       'The destination stream errored while writing data.',

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

@@ -299,7 +299,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
       );
     });
 
-    controller.abort();
+    await serverAct(() => {
+      controller.abort();
+    });
 
     const result = await resultPromise;
 
@@ -329,7 +331,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
     await jest.runAllTimers();
 
     const theReason = new Error('aborted for reasons');
-    controller.abort(theReason);
+    await serverAct(() => {
+      controller.abort(theReason);
+    });
 
     let rejected = false;
     let prelude;
@@ -447,7 +451,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
       });
     });
 
-    controller.abort('foobar');
+    await serverAct(() => {
+      controller.abort('foobar');
+    });
 
     await resultPromise;
 
@@ -489,13 +495,63 @@ describe('ReactDOMFizzStaticBrowser', () => {
       });
     });
 
-    controller.abort(new Error('uh oh'));
+    await serverAct(() => {
+      controller.abort(new Error('uh oh'));
+    });
 
     await resultPromise;
 
     expect(errors).toEqual(['uh oh', 'uh oh']);
   });
 
+  it('currently uses the abort reason when an abort listener synchronously rejects pending work', async () => {
+    let reject;
+    const rejectedPromise = new Promise((resolve, rejectPromise) => {
+      reject = rejectPromise;
+    });
+    const haltedPromise = new Promise(() => {});
+    function RejectedWait() {
+      React.use(rejectedPromise);
+      return null;
+    }
+    function HaltedWait() {
+      React.use(haltedPromise);
+      return null;
+    }
+
+    const errors = [];
+    const controller = new AbortController();
+    let resultPromise;
+    await serverAct(() => {
+      resultPromise = ReactDOMFizzStatic.prerender(
+        <>
+          <Suspense fallback="Loading rejected">
+            <RejectedWait />
+          </Suspense>
+          <Suspense fallback="Loading halted">
+            <HaltedWait />
+          </Suspense>
+        </>,
+        {
+          signal: controller.signal,
+          onError(error) {
+            errors.push(error.message);
+          },
+        },
+      );
+    });
+
+    controller.signal.addEventListener('abort', () => {
+      reject(new Error('rejected during abort'));
+    });
+    await serverAct(() => {
+      controller.abort(new Error('abort reason'));
+    });
+    await resultPromise;
+
+    expect(errors).toEqual(['abort reason', 'abort reason']);
+  });
+
   it('logs an error if onHeaders throws but continues the prerender', async () => {
     const errors = [];
     function onError(error) {
@@ -562,7 +618,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
       });
     });
 
-    controller.abort();
+    await serverAct(() => {
+      controller.abort();
+    });
     const prerendered = await pendingResult;
     const postponedState = JSON.stringify(prerendered.postponed);
 
@@ -587,7 +645,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
       );
     });
 
-    controller2.abort();
+    await serverAct(() => {
+      controller2.abort();
+    });
 
     const prerendered2 = await pendingResult;
     const postponedState2 = JSON.stringify(prerendered2.postponed);
@@ -652,7 +712,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
       });
     });
 
-    controller.abort();
+    await serverAct(() => {
+      controller.abort();
+    });
 
     const prerendered = await pendingResult;
     const postponedState = JSON.stringify(prerendered.postponed);
@@ -683,7 +745,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
       );
     });
 
-    controller2.abort();
+    await serverAct(() => {
+      controller2.abort();
+    });
 
     const prerendered2 = await pendingResult;
     const postponedState2 = JSON.stringify(prerendered2.postponed);
@@ -746,7 +810,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
       });
     });
 
-    controller.abort(new Error('boom'));
+    await serverAct(() => {
+      controller.abort(new Error('boom'));
+    });
 
     const prerendered = await pendingResult;
 
@@ -808,7 +874,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
       });
     });
 
-    controller.abort();
+    await serverAct(() => {
+      controller.abort();
+    });
 
     const prerendered = await pendingResult;
 
@@ -886,7 +954,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
       });
     });
 
-    controller.abort();
+    await serverAct(() => {
+      controller.abort();
+    });
 
     const prerendered = await pendingResult;
     const postponedState = JSON.stringify(prerendered.postponed);
@@ -912,7 +982,9 @@ describe('ReactDOMFizzStaticBrowser', () => {
       );
     });
 
-    controller2.abort();
+    await serverAct(() => {
+      controller2.abort();
+    });
 
     const prerendered2 = await pendingResult;
     const postponedState2 = JSON.stringify(prerendered2.postponed);