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__/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();