[Fizz] Update postpone abort semantics when prerendering (#30541)

When aborting with a postpone value in Fizz if any tasks are still
pending in the root while prerendering the prerender will fatally error.
This is different from postponing imperatively in a root task and really
the semantics should be the same. This change updates React to treat an
abort with a postpone value as a postponed root rather than a fatal
error.

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

@@ -398,4 +398,60 @@ describe('ReactDOMFizzStatic', () => {
       </div>,
     );
   });
+
+  // @gate enablePostpone
+  it('does not fatally error when aborting with a postpone during a prerender', async () => {
+    let postponedValue;
+    try {
+      React.unstable_postpone('aborting with postpone');
+    } catch (e) {
+      postponedValue = e;
+    }
+
+    const controller = new AbortController();
+    const infinitePromise = new Promise(() => {});
+    function App() {
+      React.use(infinitePromise);
+      return <div>aborted</div>;
+    }
+
+    const pendingResult = ReactDOMFizzStatic.prerenderToNodeStream(<App />, {
+      signal: controller.signal,
+    });
+    pendingResult.catch(() => {});
+
+    await Promise.resolve();
+    controller.abort(postponedValue);
+
+    const result = await pendingResult;
+
+    await act(async () => {
+      result.prelude.pipe(writable);
+    });
+    expect(getVisibleChildren(container)).toEqual(undefined);
+  });
+
+  // @gate enablePostpone
+  it('does not fatally error when aborting with a postpone during a prerender from within', async () => {
+    let postponedValue;
+    try {
+      React.unstable_postpone('aborting with postpone');
+    } catch (e) {
+      postponedValue = e;
+    }
+
+    const controller = new AbortController();
+    function App() {
+      controller.abort(postponedValue);
+      return <div>aborted</div>;
+    }
+
+    const result = await ReactDOMFizzStatic.prerenderToNodeStream(<App />, {
+      signal: controller.signal,
+    });
+    await act(async () => {
+      result.prelude.pipe(writable);
+    });
+    expect(getVisibleChildren(container)).toEqual(undefined);
+  });
 });

packages/react-server/src/ReactFizzServer.js

@@ -3794,12 +3794,22 @@ function abortTask(task: Task, request: Request, error: mixed): void {
           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, null);
-          fatalError(request, fatal, errorInfo, null);
+          const trackedPostpones = request.trackedPostpones;
+
+          if (trackedPostpones !== null && segment !== null) {
+            // We are prerendering. We don't want to fatal when the shell postpones
+            // we just need to mark it as postponed.
+            logPostpone(request, postponeInstance.message, errorInfo, null);
+            trackPostpone(request, trackedPostpones, task, segment);
+            finishedTask(request, null, segment);
+          } else {
+            const fatal = new Error(
+              'The render was aborted with postpone when the shell is incomplete. Reason: ' +
+                postponeInstance.message,
+            );
+            logRecoverableError(request, fatal, errorInfo, null);
+            fatalError(request, fatal, errorInfo, null);
+          }
         } else {
           logRecoverableError(request, error, errorInfo, null);
           fatalError(request, error, errorInfo, null);
@@ -4102,7 +4112,7 @@ function retryRenderTask(
     task.abortSet.delete(task);
     segment.status = COMPLETED;
     finishedTask(request, task.blockedBoundary, segment);
-  } catch (thrownValue) {
+  } catch (thrownValue: mixed) {
     resetHooksState();
 
     // Reset the write pointers to where we started.
@@ -4117,7 +4127,9 @@ function retryRenderTask(
           // (unstable) API for suspending. This implementation detail can change
           // later, once we deprecate the old API in favor of `use`.
           getSuspendedThenable()
-        : thrownValue;
+        : thrownValue === AbortSigil
+          ? request.fatalError
+          : thrownValue;
 
     if (typeof x === 'object' && x !== null) {
       // $FlowFixMe[method-unbinding]
@@ -4126,7 +4138,8 @@ function retryRenderTask(
         segment.status = PENDING;
         task.thenableState = getThenableStateAfterSuspending();
         const ping = task.ping;
-        x.then(ping, ping);
+        // We've asserted that x is a thenable above
+        (x: any).then(ping, ping);
         return;
       } else if (
         enablePostpone &&
@@ -4156,25 +4169,14 @@ function retryRenderTask(
     const errorInfo = getThrownInfo(task.componentStack);
     task.abortSet.delete(task);
 
-    if (x === AbortSigil) {
-      segment.status = ABORTED;
-      erroredTask(
-        request,
-        task.blockedBoundary,
-        request.fatalError,
-        errorInfo,
-        __DEV__ && enableOwnerStacks ? task.debugTask : null,
-      );
-    } else {
-      segment.status = ERRORED;
-      erroredTask(
-        request,
-        task.blockedBoundary,
-        x,
-        errorInfo,
-        __DEV__ && enableOwnerStacks ? task.debugTask : null,
-      );
-    }
+    segment.status = ERRORED;
+    erroredTask(
+      request,
+      task.blockedBoundary,
+      x,
+      errorInfo,
+      __DEV__ && enableOwnerStacks ? task.debugTask : null,
+    );
     return;
   } finally {
     if (__DEV__) {