From b7f6c17a5dcea76cdb99b30f850ca97d3b9dbf83 Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Mon, 19 Aug 2024 19:42:33 -0700
Subject: [PATCH] [Fizz] track postpones when aborting boundaries with a
 postpone

When aborting with a postpone value boundaries are put into client rendered mode even during prerenders. This doesn't follow the postpoen semantics of the rest of fizz where during a prerender a postpone is tracked and it will leave holes in tracked postpone state that can be resumed. This change updates this behavior to match the postpones semantics between aborts and imperative postpones.
---
 .../src/__tests__/ReactDOMFizzServer-test.js       |  1 -
 packages/react-server/src/ReactFizzServer.js       | 14 +++++++++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
index 5d70a0c71ea4d..7d8707bcd3f22 100644
--- a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
+++ b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
@@ -7727,7 +7727,6 @@ describe('ReactDOMFizzServer', () => {
 
     const prerendered = await pendingPrerender;
 
-    expect(prerendered.postponed).toBe(null);
     expect(errors).toEqual([]);
     expect(postpones).toEqual(['manufactured', 'manufactured']);
 
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index 986e8673a5a2a..17c278f2b0b52 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -3857,7 +3857,6 @@ function abortTask(task: Task, request: Request, error: mixed): void {
   } else {
     boundary.pendingTasks--;
     if (boundary.status !== CLIENT_RENDERED) {
-      boundary.status = CLIENT_RENDERED;
       // We construct an errorInfo from the boundary's componentStack so the error in dev will indicate which
       // boundary the message is referring to
       const errorInfo = getThrownInfo(task.componentStack);
@@ -3870,11 +3869,24 @@ function abortTask(task: Task, request: Request, error: mixed): void {
       ) {
         const postponeInstance: Postpone = (error: any);
         logPostpone(request, postponeInstance.message, errorInfo, null);
+        if (request.trackedPostpones !== null && segment !== null) {
+          trackPostpone(request, request.trackedPostpones, task, segment);
+          finishedTask(request, task.blockedBoundary, segment);
+
+          // If this boundary was still pending then we haven't already cancelled its fallbacks.
+          // We'll need to abort the fallbacks, which will also error that parent boundary.
+          boundary.fallbackAbortableTasks.forEach(fallbackTask =>
+            abortTask(fallbackTask, request, error),
+          );
+          boundary.fallbackAbortableTasks.clear();
+          return;
+        }
         // TODO: Figure out a better signal than a magic digest value.
         errorDigest = 'POSTPONE';
       } else {
         errorDigest = logRecoverableError(request, error, errorInfo, null);
       }
+      boundary.status = CLIENT_RENDERED;
       encodeErrorForBoundary(boundary, errorDigest, error, errorInfo, true);
 
       untrackBoundary(request, boundary);