[Fizz] Destroy the stream with an error if the root throws (#20992)

* Destroy the stream with an error if the root throws

But not if the error happens inside a suspense boundary.

* Try rewriting the test to see if it works in other Node envs

Author: sebmarkbage

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

@@ -15,6 +15,7 @@ global.TextEncoder = require('util').TextEncoder;
 
 let React;
 let ReactDOMFizzServer;
+let Suspense;
 
 describe('ReactDOMFizzServer', () => {
   beforeEach(() => {
@@ -23,8 +24,18 @@ describe('ReactDOMFizzServer', () => {
     if (__EXPERIMENTAL__) {
       ReactDOMFizzServer = require('react-dom/unstable-fizz.browser');
     }
+    Suspense = React.Suspense;
   });
 
+  const theError = new Error('This is an error');
+  function Throw() {
+    throw theError;
+  }
+  const theInfinitePromise = new Promise(() => {});
+  function InfiniteSuspend() {
+    throw theInfinitePromise;
+  }
+
   async function readResult(stream) {
     const reader = stream.getReader();
     let result = '';
@@ -45,4 +56,58 @@ describe('ReactDOMFizzServer', () => {
     const result = await readResult(stream);
     expect(result).toBe('<div>hello world</div>');
   });
+
+  // @gate experimental
+  it('should error the stream when an error is thrown at the root', async () => {
+    const stream = ReactDOMFizzServer.renderToReadableStream(
+      <div>
+        <Throw />
+      </div>,
+    );
+
+    let caughtError = null;
+    let result = '';
+    try {
+      result = await readResult(stream);
+    } catch (x) {
+      caughtError = x;
+    }
+    expect(caughtError).toBe(theError);
+    expect(result).toBe('');
+  });
+
+  // @gate experimental
+  it('should error the stream when an error is thrown inside a fallback', async () => {
+    const stream = ReactDOMFizzServer.renderToReadableStream(
+      <div>
+        <Suspense fallback={<Throw />}>
+          <InfiniteSuspend />
+        </Suspense>
+      </div>,
+    );
+
+    let caughtError = null;
+    let result = '';
+    try {
+      result = await readResult(stream);
+    } catch (x) {
+      caughtError = x;
+    }
+    expect(caughtError).toBe(theError);
+    expect(result).toBe('');
+  });
+
+  // @gate experimental
+  it('should not error the stream when an error is thrown inside suspense boundary', async () => {
+    const stream = ReactDOMFizzServer.renderToReadableStream(
+      <div>
+        <Suspense fallback={<div>Loading</div>}>
+          <Throw />
+        </Suspense>
+      </div>,
+    );
+
+    const result = await readResult(stream);
+    expect(result).toContain('Loading');
+  });
 });

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

@@ -13,6 +13,7 @@
 let Stream;
 let React;
 let ReactDOMFizzServer;
+let Suspense;
 
 describe('ReactDOMFizzServer', () => {
   beforeEach(() => {
@@ -22,21 +23,96 @@ describe('ReactDOMFizzServer', () => {
       ReactDOMFizzServer = require('react-dom/unstable-fizz');
     }
     Stream = require('stream');
+    Suspense = React.Suspense;
   });
 
   function getTestWritable() {
     const writable = new Stream.PassThrough();
     writable.setEncoding('utf8');
-    writable.result = '';
-    writable.on('data', chunk => (writable.result += chunk));
-    return writable;
+    const output = {result: '', error: undefined};
+    writable.on('data', chunk => {
+      output.result += chunk;
+    });
+    writable.on('error', error => {
+      output.error = error;
+    });
+    const completed = new Promise(resolve => {
+      writable.on('finish', () => {
+        resolve();
+      });
+      writable.on('error', () => {
+        resolve();
+      });
+    });
+    return {writable, completed, output};
+  }
+
+  const theError = new Error('This is an error');
+  function Throw() {
+    throw theError;
+  }
+  const theInfinitePromise = new Promise(() => {});
+  function InfiniteSuspend() {
+    throw theInfinitePromise;
   }
 
   // @gate experimental
   it('should call pipeToNodeWritable', () => {
-    const writable = getTestWritable();
+    const {writable, output} = getTestWritable();
     ReactDOMFizzServer.pipeToNodeWritable(<div>hello world</div>, writable);
     jest.runAllTimers();
-    expect(writable.result).toBe('<div>hello world</div>');
+    expect(output.result).toBe('<div>hello world</div>');
+  });
+
+  // @gate experimental
+  it('should error the stream when an error is thrown at the root', async () => {
+    const {writable, output, completed} = getTestWritable();
+    ReactDOMFizzServer.pipeToNodeWritable(
+      <div>
+        <Throw />
+      </div>,
+      writable,
+    );
+
+    await completed;
+
+    expect(output.error).toBe(theError);
+    expect(output.result).toBe('');
+  });
+
+  // @gate experimental
+  it('should error the stream when an error is thrown inside a fallback', async () => {
+    const {writable, output, completed} = getTestWritable();
+    ReactDOMFizzServer.pipeToNodeWritable(
+      <div>
+        <Suspense fallback={<Throw />}>
+          <InfiniteSuspend />
+        </Suspense>
+      </div>,
+      writable,
+    );
+
+    await completed;
+
+    expect(output.error).toBe(theError);
+    expect(output.result).toBe('');
+  });
+
+  // @gate experimental
+  it('should not error the stream when an error is thrown inside suspense boundary', async () => {
+    const {writable, output, completed} = getTestWritable();
+    ReactDOMFizzServer.pipeToNodeWritable(
+      <div>
+        <Suspense fallback={<div>Loading</div>}>
+          <Throw />
+        </Suspense>
+      </div>,
+      writable,
+    );
+
+    await completed;
+
+    expect(output.error).toBe(undefined);
+    expect(output.result).toContain('Loading');
   });
 });

packages/react-noop-renderer/src/ReactNoopFlightServer.js

@@ -32,6 +32,7 @@ const ReactNoopFlightServer = ReactFlightServer({
   },
   completeWriting(destination: Destination): void {},
   close(destination: Destination): void {},
+  closeWithError(destination: Destination, error: mixed): void {},
   flushBuffered(destination: Destination): void {},
   convertStringToBuffer(content: string): Uint8Array {
     return Buffer.from(content, 'utf8');

packages/react-noop-renderer/src/ReactNoopServer.js

@@ -74,6 +74,7 @@ const ReactNoopServer = ReactFizzServer({
   },
   completeWriting(destination: Destination): void {},
   close(destination: Destination): void {},
+  closeWithError(destination: Destination, error: mixed): void {},
   flushBuffered(destination: Destination): void {},
 
   createResponseState(): null {

packages/react-server/src/ReactFizzServer.js

@@ -22,6 +22,7 @@ import {
   completeWriting,
   flushBuffered,
   close,
+  closeWithError,
 } from './ReactServerStreamConfig';
 import {
   writePlaceholder,
@@ -205,7 +206,7 @@ function fatalError(request: Request, error: mixed): void {
   // a suspense boundary or if the root suspense boundary's fallback errors.
   // It's also called if React itself or its host configs errors.
   request.status = CLOSED;
-  // TODO: Destroy the stream with an error. We weren't able to complete the root.
+  closeWithError(request.destination, error);
 }
 
 function renderNode(
@@ -237,6 +238,7 @@ function renderNode(
         // Something suspended, we'll need to create a new segment and resolve it later.
         const insertionIndex = segment.chunks.length;
         const newSegment = createPendingSegment(request, insertionIndex, null);
+        segment.children.push(newSegment);
         const suspendedWork = createSuspendedWork(
           request,
           node,
@@ -273,6 +275,7 @@ function renderNode(
       insertionIndex,
       newBoundary,
     );
+    segment.children.push(boundarySegment);
     // We create suspended work for the fallback because we don't want to actually work
     // on it yet in case we finish the main content, so we queue for later.
     const suspendedFallbackWork = createSuspendedWork(
@@ -326,6 +329,7 @@ function errorWork(
     fatalError(request, error);
   } else if (!boundary.forceClientRender) {
     boundary.forceClientRender = true;
+
     // Regardless of what happens next, this boundary won't be displayed,
     // so we can flush it, if the parent already flushed.
     if (boundary.parentFlushed) {

packages/react-server/src/ReactServerStreamConfigBrowser.js

@@ -39,3 +39,18 @@ const textEncoder = new TextEncoder();
 export function convertStringToBuffer(content: string): Uint8Array {
   return textEncoder.encode(content);
 }
+
+export function closeWithError(destination: Destination, error: mixed): void {
+  if (typeof destination.error === 'function') {
+    // $FlowFixMe: This is an Error object or the destination accepts other types.
+    destination.error(error);
+  } else {
+    // Earlier implementations doesn't support this method. In that environment you're
+    // supposed to throw from a promise returned but we don't return a promise in our
+    // approach. We could fork this implementation but this is environment is an edge
+    // case to begin with. It's even less common to run this in an older environment.
+    // Even then, this is not where errors are supposed to happen and they get reported
+    // to a global callback in addition to this anyway. So it's fine just to close this.
+    destination.close();
+  }
+}

packages/react-server/src/ReactServerStreamConfigNode.js

@@ -64,3 +64,8 @@ export function close(destination: Destination) {
 export function convertStringToBuffer(content: string): Uint8Array {
   return Buffer.from(content, 'utf8');
 }
+
+export function closeWithError(destination: Destination, error: mixed): void {
+  // $FlowFixMe: This is an Error object or the destination accepts other types.
+  destination.destroy(error);
+}

packages/react-server/src/forks/ReactServerStreamConfig.custom.js

@@ -32,4 +32,5 @@ export const writeChunk = $$$hostConfig.writeChunk;
 export const completeWriting = $$$hostConfig.completeWriting;
 export const flushBuffered = $$$hostConfig.flushBuffered;
 export const close = $$$hostConfig.close;
+export const closeWithError = $$$hostConfig.closeWithError;
 export const convertStringToBuffer = $$$hostConfig.convertStringToBuffer;