[Fizz] Fix reentrancy bug

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

@@ -353,13 +353,16 @@ describe('ReactDOMFizzServer', () => {
 
     await act(async () => {
       const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
-        <Suspense fallback={<Text text="Loading A..." />}>
-          <>
-            <Text text="This will show A: " />
-            <div>
-              <AsyncText text="A" />
-            </div>
-          </>
+        // We use two nested boundaries to flush out coverage of an old reentrancy bug.
+        <Suspense fallback="Loading...">
+          <Suspense fallback={<Text text="Loading A..." />}>
+            <>
+              <Text text="This will show A: " />
+              <div>
+                <AsyncText text="A" />
+              </div>
+            </>
+          </Suspense>
         </Suspense>,
         writableA,
         {

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

@@ -99,7 +99,7 @@ describe('ReactDOMFizzServer', () => {
       }
       return 'Done';
     }
-    let isComplete = false;
+    let isCompleteCalls = 0;
     const {writable, output} = getTestWritable();
     const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
       <div>
@@ -110,21 +110,21 @@ describe('ReactDOMFizzServer', () => {
       writable,
       {
         onCompleteAll() {
-          isComplete = true;
+          isCompleteCalls++;
         },
       },
     );
     await jest.runAllTimers();
     expect(output.result).toBe('');
-    expect(isComplete).toBe(false);
+    expect(isCompleteCalls).toBe(0);
     // Resolve the loading.
     hasLoaded = true;
     await resolve();
 
     await jest.runAllTimers();
 
     expect(output.result).toBe('');
-    expect(isComplete).toBe(true);
+    expect(isCompleteCalls).toBe(1);
 
     // First we write our header.
     output.result +=
@@ -244,6 +244,7 @@ describe('ReactDOMFizzServer', () => {
 
   // @gate experimental
   it('should be able to complete by aborting even if the promise never resolves', async () => {
+    let isCompleteCalls = 0;
     const {writable, output, completed} = getTestWritable();
     const {startWriting, abort} = ReactDOMFizzServer.pipeToNodeWritable(
       <div>
@@ -252,18 +253,60 @@ describe('ReactDOMFizzServer', () => {
         </Suspense>
       </div>,
       writable,
+      {
+        onCompleteAll() {
+          isCompleteCalls++;
+        },
+      },
+    );
+    startWriting();
+
+    jest.runAllTimers();
+
+    expect(output.result).toContain('Loading');
+    expect(isCompleteCalls).toBe(0);
+
+    abort();
+
+    await completed;
+
+    expect(output.error).toBe(undefined);
+    expect(output.result).toContain('Loading');
+    expect(isCompleteCalls).toBe(1);
+  });
+
+  // @gate experimental
+  it('should be able to complete by abort when the fallback is also suspended', async () => {
+    let isCompleteCalls = 0;
+    const {writable, output, completed} = getTestWritable();
+    const {startWriting, abort} = ReactDOMFizzServer.pipeToNodeWritable(
+      <div>
+        <Suspense fallback="Loading">
+          <Suspense fallback={<InfiniteSuspend />}>
+            <InfiniteSuspend />
+          </Suspense>
+        </Suspense>
+      </div>,
+      writable,
+      {
+        onCompleteAll() {
+          isCompleteCalls++;
+        },
+      },
     );
     startWriting();
 
     jest.runAllTimers();
 
     expect(output.result).toContain('Loading');
+    expect(isCompleteCalls).toBe(0);
 
     abort();
 
     await completed;
 
     expect(output.error).toBe(undefined);
     expect(output.result).toContain('Loading');
+    expect(isCompleteCalls).toBe(1);
   });
 });

packages/react-server/src/ReactFizzServer.js

@@ -1006,8 +1006,8 @@ function abortTask(task: Task): void {
   const segment = task.blockedSegment;
   segment.status = ABORTED;
 
-  request.allPendingTasks--;
   if (boundary === null) {
+    request.allPendingTasks--;
     // We didn't complete the root so we have nothing to show. We can close
     // the request;
     if (request.status !== CLOSED) {
@@ -1017,18 +1017,23 @@ function abortTask(task: Task): void {
   } else {
     boundary.pendingTasks--;
 
-    // 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(abortTask, request);
-    boundary.fallbackAbortableTasks.clear();
-
-    if (!boundary.forceClientRender) {
-      boundary.forceClientRender = true;
-      if (boundary.parentFlushed) {
-        request.clientRenderedBoundaries.push(boundary);
+    if (boundary.fallbackAbortableTasks.size > 0) {
+      // 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.
+      // This means that we don't have to client render this boundary because its parent
+      // will be client rendered anyway.
+      boundary.fallbackAbortableTasks.forEach(abortTask, request);
+      boundary.fallbackAbortableTasks.clear();
+    } else {
+      if (!boundary.forceClientRender) {
+        boundary.forceClientRender = true;
+        if (boundary.parentFlushed) {
+          request.clientRenderedBoundaries.push(boundary);
+        }
       }
     }
 
+    request.allPendingTasks--;
     if (request.allPendingTasks === 0) {
       const onCompleteAll = request.onCompleteAll;
       onCompleteAll();
@@ -1060,9 +1065,6 @@ function finishedTask(
       // This already errored.
     } else if (boundary.pendingTasks === 0) {
       // This must have been the last segment we were waiting on. This boundary is now complete.
-      // We can now cancel any pending task on the fallback since we won't need to show it anymore.
-      boundary.fallbackAbortableTasks.forEach(abortTaskSoft, request);
-      boundary.fallbackAbortableTasks.clear();
       if (segment.parentFlushed) {
         // Our parent segment already flushed, so we need to schedule this segment to be emitted.
         boundary.completedSegments.push(segment);
@@ -1072,6 +1074,11 @@ function finishedTask(
         // parent flushed, we need to schedule the boundary to be emitted.
         request.completedBoundaries.push(boundary);
       }
+      // We can now cancel any pending task on the fallback since we won't need to show it anymore.
+      // This needs to happen after we read the parentFlushed flags because aborting can finish
+      // work which can trigger user code, which can start flushing, which can change those flags.
+      boundary.fallbackAbortableTasks.forEach(abortTaskSoft, request);
+      boundary.fallbackAbortableTasks.clear();
     } else {
       if (segment.parentFlushed) {
         // Our parent already flushed, so we need to schedule this segment to be emitted.