Move onCompleteAll to .allReady Promise (#24025)

* Move onCompleteAll to .allReady Promise

The onCompleteAll callback can sometimes resolve before the promise that
returns the stream which is tough to coordinate. A more idiomatic API
for a one shot event is a Promise.

That way the way you render for SEO or SSG is:

const stream = await renderToReadableStream(...);
await stream.readyAll;
respondWith(stream);

Ideally this should be a sub-class of ReadableStream but we don't yet
compile these to ES6 and they'd had to be to native class to subclass
a native stream.

I have other ideas for overriding the .tee() method in a subclass anyway.
So this is inline with that strategy.

* Reject the Promise on fatal errors

Author: sebmarkbage

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

@@ -105,12 +105,10 @@ describe('ReactDOMFizzServer', () => {
           <Wait />
         </Suspense>
       </div>,
-      {
-        onCompleteAll() {
-          isComplete = true;
-        },
-      },
     );
+
+    stream.allReady.then(() => (isComplete = true));
+
     await jest.runAllTimers();
     expect(isComplete).toBe(false);
     // Resolve the loading.

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

@@ -32,23 +32,36 @@ type Options = {|
   bootstrapModules?: Array<string>,
   progressiveChunkSize?: number,
   signal?: AbortSignal,
-  onCompleteAll?: () => void,
   onError?: (error: mixed) => void,
 |};
 
+// TODO: Move to sub-classing ReadableStream.
+type ReactDOMServerReadableStream = ReadableStream & {
+  allReady: Promise<void>,
+};
+
 function renderToReadableStream(
   children: ReactNodeList,
   options?: Options,
-): Promise<ReadableStream> {
+): Promise<ReactDOMServerReadableStream> {
   return new Promise((resolve, reject) => {
+    let onFatalError;
+    let onCompleteAll;
+    const allReady = new Promise((res, rej) => {
+      onCompleteAll = res;
+      onFatalError = rej;
+    });
+
     function onCompleteShell() {
-      const stream = new ReadableStream({
+      const stream: ReactDOMServerReadableStream = (new ReadableStream({
         type: 'bytes',
         pull(controller) {
           startFlowing(request, controller);
         },
         cancel(reason) {},
-      });
+      }): any);
+      // TODO: Move to sub-classing ReadableStream.
+      stream.allReady = allReady;
       resolve(stream);
     }
     function onErrorShell(error: mixed) {
@@ -66,9 +79,10 @@ function renderToReadableStream(
       createRootFormatContext(options ? options.namespaceURI : undefined),
       options ? options.progressiveChunkSize : undefined,
       options ? options.onError : undefined,
-      options ? options.onCompleteAll : undefined,
+      onCompleteAll,
       onCompleteShell,
       onErrorShell,
+      onFatalError,
     );
     if (options && options.signal) {
       const signal = options.signal;

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

@@ -65,6 +65,7 @@ function createRequestImpl(children: ReactNodeList, options: void | Options) {
     options ? options.onCompleteAll : undefined,
     options ? options.onCompleteShell : undefined,
     options ? options.onErrorShell : undefined,
+    undefined,
   );
 }
 

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

@@ -67,6 +67,8 @@ function renderToStringImpl(
     onError,
     undefined,
     onCompleteShell,
+    undefined,
+    undefined,
   );
   startWork(request);
   // If anything suspended and is still pending, we'll abort it before writing.

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

@@ -83,6 +83,7 @@ function renderToNodeStreamImpl(
     onError,
     onCompleteAll,
     undefined,
+    undefined,
   );
   destination.request = request;
   startWork(request);

packages/react-server/src/ReactFizzServer.js

@@ -204,6 +204,7 @@ export opaque type Request = {
   // onErrorShell is called when the shell didn't complete. That means you probably want to
   // emit a different response to the stream instead.
   onErrorShell: (error: mixed) => void,
+  onFatalError: (error: mixed) => void,
 };
 
 // This is a default heuristic for how to split up the HTML content into progressive
@@ -238,6 +239,7 @@ export function createRequest(
   onCompleteAll: void | (() => void),
   onCompleteShell: void | (() => void),
   onErrorShell: void | ((error: mixed) => void),
+  onFatalError: void | ((error: mixed) => void),
 ): Request {
   const pingedTasks = [];
   const abortSet: Set<Task> = new Set();
@@ -263,6 +265,7 @@ export function createRequest(
     onCompleteAll: onCompleteAll === undefined ? noop : onCompleteAll,
     onCompleteShell: onCompleteShell === undefined ? noop : onCompleteShell,
     onErrorShell: onErrorShell === undefined ? noop : onErrorShell,
+    onFatalError: onFatalError === undefined ? noop : onFatalError,
   };
   // This segment represents the root fallback.
   const rootSegment = createPendingSegment(request, 0, null, rootFormatContext);