From dbd968be6d070aa3a366ae62b736777563ef27cf Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Thu, 3 Mar 2022 12:15:26 -0500
Subject: [PATCH] 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.
---
 .../ReactDOMFizzServerBrowser-test.js         |  8 +++-----
 .../src/server/ReactDOMFizzServerBrowser.js   | 19 +++++++++++++++----
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js
index 87d04bdd07655..93e30aa2305be 100644
--- a/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js
+++ b/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.
diff --git a/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js b/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js
index a07967ca34f98..768c73640710c 100644
--- a/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js
+++ b/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js
@@ -32,23 +32,34 @@ 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 onCompleteAll;
+    const allReady = new Promise(resolve => {
+      onCompleteAll = resolve;
+    });
+
     function onCompleteShell() {
-      const stream = new ReadableStream({
+      const stream: ReactDOMServerReadableStream = new ReadableStream({
         type: 'bytes',
         pull(controller) {
           startFlowing(request, controller);
         },
         cancel(reason) {},
       });
+      // TODO: Move to sub-classing ReadableStream.
+      stream.allReady = allReady;
       resolve(stream);
     }
     function onErrorShell(error: mixed) {
@@ -66,7 +77,7 @@ function renderToReadableStream(
       createRootFormatContext(options ? options.namespaceURI : undefined),
       options ? options.progressiveChunkSize : undefined,
       options ? options.onError : undefined,
-      options ? options.onCompleteAll : undefined,
+      onCompleteAll,
       onCompleteShell,
       onErrorShell,
     );