[Fizz] Expose callbacks in options for when various stages of the content is done

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

@@ -336,7 +336,6 @@ describe('ReactDOMFizzServer', () => {
       writable.write(chunk, encoding, next);
     };
 
-    writable.write('<div id="container-A">');
     await act(async () => {
       const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
         <Suspense fallback={<Text text="Loading A..." />}>
@@ -346,13 +345,17 @@ describe('ReactDOMFizzServer', () => {
           </div>
         </Suspense>,
         writableA,
-        {identifierPrefix: 'A_'},
+        {
+          identifierPrefix: 'A_',
+          onReadyToStream() {
+            writableA.write('<div id="container-A">');
+            startWriting();
+            writableA.write('</div>');
+          },
+        },
       );
-      startWriting();
     });
-    writable.write('</div>');
 
-    writable.write('<div id="container-B">');
     await act(async () => {
       const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
         <Suspense fallback={<Text text="Loading B..." />}>
@@ -362,11 +365,16 @@ describe('ReactDOMFizzServer', () => {
           </div>
         </Suspense>,
         writableB,
-        {identifierPrefix: 'B_'},
+        {
+          identifierPrefix: 'B_',
+          onReadyToStream() {
+            writableB.write('<div id="container-B">');
+            startWriting();
+            writableB.write('</div>');
+          },
+        },
       );
-      startWriting();
     });
-    writable.write('</div>');
 
     expect(getVisibleChildren(container)).toEqual([
       <div id="container-A">Loading A...</div>,

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

@@ -58,12 +58,56 @@ describe('ReactDOMFizzServer', () => {
     expect(result).toBe('<div>hello world</div>');
   });
 
+  // @gate experimental
+  it('emits all HTML as one unit if we wait until the end to start', async () => {
+    let hasLoaded = false;
+    let resolve;
+    const promise = new Promise(r => (resolve = r));
+    function Wait() {
+      if (!hasLoaded) {
+        throw promise;
+      }
+      return 'Done';
+    }
+    let isComplete = false;
+    const stream = ReactDOMFizzServer.renderToReadableStream(
+      <div>
+        <Suspense fallback="Loading">
+          <Wait />
+        </Suspense>
+      </div>,
+      {
+        onCompleteAll() {
+          isComplete = true;
+        },
+      },
+    );
+    await jest.runAllTimers();
+    expect(isComplete).toBe(false);
+    // Resolve the loading.
+    hasLoaded = true;
+    await resolve();
+
+    await jest.runAllTimers();
+
+    expect(isComplete).toBe(true);
+
+    const result = await readResult(stream);
+    expect(result).toBe('<div><!--$-->Done<!--/$--></div>');
+  });
+
   // @gate experimental
   it('should error the stream when an error is thrown at the root', async () => {
+    const reportedErrors = [];
     const stream = ReactDOMFizzServer.renderToReadableStream(
       <div>
         <Throw />
       </div>,
+      {
+        onError(x) {
+          reportedErrors.push(x);
+        },
+      },
     );
 
     let caughtError = null;
@@ -75,16 +119,23 @@ describe('ReactDOMFizzServer', () => {
     }
     expect(caughtError).toBe(theError);
     expect(result).toBe('');
+    expect(reportedErrors).toEqual([theError]);
   });
 
   // @gate experimental
   it('should error the stream when an error is thrown inside a fallback', async () => {
+    const reportedErrors = [];
     const stream = ReactDOMFizzServer.renderToReadableStream(
       <div>
         <Suspense fallback={<Throw />}>
           <InfiniteSuspend />
         </Suspense>
       </div>,
+      {
+        onError(x) {
+          reportedErrors.push(x);
+        },
+      },
     );
 
     let caughtError = null;
@@ -96,20 +147,28 @@ describe('ReactDOMFizzServer', () => {
     }
     expect(caughtError).toBe(theError);
     expect(result).toBe('');
+    expect(reportedErrors).toEqual([theError]);
   });
 
   // @gate experimental
   it('should not error the stream when an error is thrown inside suspense boundary', async () => {
+    const reportedErrors = [];
     const stream = ReactDOMFizzServer.renderToReadableStream(
       <div>
         <Suspense fallback={<div>Loading</div>}>
           <Throw />
         </Suspense>
       </div>,
+      {
+        onError(x) {
+          reportedErrors.push(x);
+        },
+      },
     );
 
     const result = await readResult(stream);
     expect(result).toContain('Loading');
+    expect(reportedErrors).toEqual([theError]);
   });
 
   // @gate experimental

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

@@ -86,14 +86,68 @@ describe('ReactDOMFizzServer', () => {
     );
   });
 
+  // @gate experimental
+  it('emits all HTML as one unit if we wait until the end to start', async () => {
+    let hasLoaded = false;
+    let resolve;
+    const promise = new Promise(r => (resolve = r));
+    function Wait() {
+      if (!hasLoaded) {
+        throw promise;
+      }
+      return 'Done';
+    }
+    let isComplete = false;
+    const {writable, output} = getTestWritable();
+    const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
+      <div>
+        <Suspense fallback="Loading">
+          <Wait />
+        </Suspense>
+      </div>,
+      writable,
+      {
+        onCompleteAll() {
+          isComplete = true;
+        },
+      },
+    );
+    await jest.runAllTimers();
+    expect(output.result).toBe('');
+    expect(isComplete).toBe(false);
+    // Resolve the loading.
+    hasLoaded = true;
+    await resolve();
+
+    await jest.runAllTimers();
+
+    expect(output.result).toBe('');
+    expect(isComplete).toBe(true);
+
+    // First we write our header.
+    output.result +=
+      '<!doctype html><html><head><title>test</title><head><body>';
+    // Then React starts writing.
+    startWriting();
+    expect(output.result).toBe(
+      '<!doctype html><html><head><title>test</title><head><body><div><!--$-->Done<!--/$--></div>',
+    );
+  });
+
   // @gate experimental
   it('should error the stream when an error is thrown at the root', async () => {
+    const reportedErrors = [];
     const {writable, output, completed} = getTestWritable();
     ReactDOMFizzServer.pipeToNodeWritable(
       <div>
         <Throw />
       </div>,
       writable,
+      {
+        onError(x) {
+          reportedErrors.push(x);
+        },
+      },
     );
 
     // The stream is errored even if we haven't started writing.
@@ -102,10 +156,13 @@ describe('ReactDOMFizzServer', () => {
 
     expect(output.error).toBe(theError);
     expect(output.result).toBe('');
+    // This type of error is reported to the error callback too.
+    expect(reportedErrors).toEqual([theError]);
   });
 
   // @gate experimental
   it('should error the stream when an error is thrown inside a fallback', async () => {
+    const reportedErrors = [];
     const {writable, output, completed} = getTestWritable();
     const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
       <div>
@@ -114,17 +171,24 @@ describe('ReactDOMFizzServer', () => {
         </Suspense>
       </div>,
       writable,
+      {
+        onError(x) {
+          reportedErrors.push(x);
+        },
+      },
     );
     startWriting();
 
     await completed;
 
     expect(output.error).toBe(theError);
     expect(output.result).toBe('');
+    expect(reportedErrors).toEqual([theError]);
   });
 
   // @gate experimental
   it('should not error the stream when an error is thrown inside suspense boundary', async () => {
+    const reportedErrors = [];
     const {writable, output, completed} = getTestWritable();
     const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
       <div>
@@ -133,13 +197,20 @@ describe('ReactDOMFizzServer', () => {
         </Suspense>
       </div>,
       writable,
+      {
+        onError(x) {
+          reportedErrors.push(x);
+        },
+      },
     );
     startWriting();
 
     await completed;
 
     expect(output.error).toBe(undefined);
     expect(output.result).toContain('Loading');
+    // While no error is reported to the stream, the error is reported to the callback.
+    expect(reportedErrors).toEqual([theError]);
   });
 
   // @gate experimental

packages/react-dom/src/server/ReactDOMFizzServerBrowser.js

@@ -22,6 +22,9 @@ type Options = {
   identifierPrefix?: string,
   progressiveChunkSize?: number,
   signal?: AbortSignal,
+  onReadyToStream?: () => void,
+  onCompleteAll?: () => void,
+  onError?: (error: mixed) => void,
 };
 
 function renderToReadableStream(
@@ -37,21 +40,31 @@ function renderToReadableStream(
     };
     signal.addEventListener('abort', listener);
   }
-  return new ReadableStream({
+  const stream = new ReadableStream({
     start(controller) {
       request = createRequest(
         children,
         controller,
         createResponseState(options ? options.identifierPrefix : undefined),
         options ? options.progressiveChunkSize : undefined,
+        options ? options.onError : undefined,
+        options ? options.onCompleteAll : undefined,
+        options ? options.onReadyToStream : undefined,
       );
       startWork(request);
     },
     pull(controller) {
-      startFlowing(request);
+      // Pull is called immediately even if the stream is not passed to anything.
+      // That's buffering too early. We want to start buffering once the stream
+      // is actually used by something so we can give it the best result possible
+      // at that point.
+      if (stream.locked) {
+        startFlowing(request);
+      }
     },
     cancel(reason) {},
   });
+  return stream;
 }
 
 export {renderToReadableStream};

packages/react-dom/src/server/ReactDOMFizzServerNode.js

@@ -26,6 +26,9 @@ function createDrainHandler(destination, request) {
 type Options = {
   identifierPrefix?: string,
   progressiveChunkSize?: number,
+  onReadyToStream?: () => void,
+  onCompleteAll?: () => void,
+  onError?: (error: mixed) => void,
 };
 
 type Controls = {
@@ -44,6 +47,9 @@ function pipeToNodeWritable(
     destination,
     createResponseState(options ? options.identifierPrefix : undefined),
     options ? options.progressiveChunkSize : undefined,
+    options ? options.onError : undefined,
+    options ? options.onCompleteAll : undefined,
+    options ? options.onReadyToStream : undefined,
   );
   let hasStartedFlowing = false;
   startWork(request);

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

@@ -217,6 +217,9 @@ const ReactNoopServer = ReactFizzServer({
 
 type Options = {
   progressiveChunkSize?: number,
+  onReadyToStream?: () => void,
+  onCompleteAll?: () => void,
+  onError?: (error: mixed) => void,
 };
 
 function render(children: React$Element<any>, options?: Options): Destination {
@@ -234,6 +237,9 @@ function render(children: React$Element<any>, options?: Options): Destination {
     destination,
     null,
     options ? options.progressiveChunkSize : undefined,
+    options ? options.onError : undefined,
+    options ? options.onCompleteAll : undefined,
+    options ? options.onReadyToStream : undefined,
   );
   ReactNoopServer.startWork(request);
   ReactNoopServer.startFlowing(request);