Aborting a replay task at the root should still complete the shell

sebmarkbage · sebmarkbage · commit 104995f98abf · 2023-10-04T00:35:28.000-04:00
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js
@@ -1065,4 +1065,65 @@ describe('ReactDOMFizzStaticBrowser', () => {
       <div>Loading...</div>,
     ]);
   });
+
+  // @gate enablePostpone
+  it('can abort the resume', async () => {
+    let prerendering = true;
+    const infinitePromise = new Promise(() => {});
+    function Postpone() {
+      if (prerendering) {
+        React.unstable_postpone();
+      }
+      return 'Hello';
+    }
+
+    function App() {
+      if (!prerendering) {
+        React.use(infinitePromise);
+      }
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <Postpone />
+          </Suspense>
+        </div>
+      );
+    }
+
+    const prerendered = await ReactDOMFizzStatic.prerender(<App />);
+    expect(prerendered.postponed).not.toBe(null);
+
+    await readIntoContainer(prerendered.prelude);
+
+    expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
+
+    prerendering = false;
+
+    const controller = new AbortController();
+
+    const errors = [];
+
+    const resumedPromise = ReactDOMFizzServer.resume(
+      <App />,
+      JSON.parse(JSON.stringify(prerendered.postponed)),
+      {
+        signal: controller.signal,
+        onError(x) {
+          errors.push(x);
+        },
+      },
+    );
+
+    controller.abort('abort');
+
+    const resumed = await resumedPromise;
+    await resumed.allReady;
+
+    expect(errors).toEqual(['abort']);
+
+    await readIntoContainer(resumed);
+
+    // Client rendered
+    expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
+  });
 });
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
@@ -3088,14 +3088,14 @@ function abortTask(task: Task, request: Request, error: mixed): void {
   }
 
   if (boundary === null) {
-    request.allPendingTasks--;
     if (request.status !== CLOSING && request.status !== CLOSED) {
       const replay: null | ReplaySet = task.replay;
       if (replay === null) {
         // We didn't complete the root so we have nothing to show. We can close
         // the request;
         logRecoverableError(request, error);
         fatalError(request, error);
+        return;
       } else {
         // If the shell aborts during a replay, that's not a fatal error. Instead
         // we should be able to recover by client rendering all the root boundaries in
@@ -3112,6 +3112,12 @@ function abortTask(task: Task, request: Request, error: mixed): void {
             errorDigest,
           );
         }
+        request.pendingRootTasks--;
+        if (request.pendingRootTasks === 0) {
+          request.onShellError = noop;
+          const onShellReady = request.onShellReady;
+          onShellReady();
+        }
       }
     }
   } else {
@@ -3148,12 +3154,12 @@ function abortTask(task: Task, request: Request, error: mixed): void {
       abortTask(fallbackTask, request, error),
     );
     boundary.fallbackAbortableTasks.clear();
+  }
 
-    request.allPendingTasks--;
-    if (request.allPendingTasks === 0) {
-      const onAllReady = request.onAllReady;
-      onAllReady();
-    }
+  request.allPendingTasks--;
+  if (request.allPendingTasks === 0) {
+    const onAllReady = request.onAllReady;
+    onAllReady();
   }
 }
 
