[Fizz] Fix reentrancy bug (#21270)

* Fix reentrancy bug

* Fix another reentrancy bug

There's also an issue if we try to schedule something to be client
rendered if its fallback hasn't rendered yet. So we don't do it
in that case.

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

@@ -1172,8 +1172,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) {
@@ -1183,18 +1183,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();
@@ -1226,9 +1231,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);
@@ -1238,6 +1240,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.