Merge pull request #41 from zth/feature/stream-simplification

Translate stream asset containers to callbacks

PreloadInsertingStream.mjs

@@ -0,0 +1,90 @@
+import { Writable } from "stream";
+
+function asRelayDataAppend(relayData) {
+  return `window.__RELAY_DATA.push(${JSON.stringify(relayData)})`;
+}
+
+export default class PreloadInsertingStream extends Writable {
+  constructor(writable) {
+    super();
+    this._queryData = [];
+    this._assetLoaderTags = [];
+    this._writable = writable;
+  }
+
+  /**
+   * Should be invoked when a new Relay query is started or updated.
+   *
+   * @param {string} id
+   *   The query ID used to track the request between server and client.
+   * @param {*} response
+   *   The Relay response or undefined in case this request is just being initiated.
+   * @param {boolean|undefined} final
+   *   Whether this is the final response for this query (or undefined in case the query is just being initiated).
+   */
+  onQuery(id, response, final) {
+    this._queryData.push({ id, response, final })
+  }
+
+  /**
+   * Should be invoked when a new asset has been used.
+   *
+   * @param {string} loaderTag
+   *   The HTML tag (e.g. `link` or `script`) that loads the required resource.
+   */
+  onAssetPreload(loaderTag) {
+    this._assetLoaderTags.push(loaderTag)
+  }
+
+  /**
+   * Generates the HTML tags needed to load assets used since the last call.
+   *
+   * @returns {string}
+   *   An HTML string containing Relay Data snippets and asset loading tags.
+   */
+  _generateNewScriptTagsSinceLastCall() {
+    let scriptTags = '';
+
+    if (this._queryData.length > 0) {
+      scriptTags += `<script type="text/javascript" class="__relay_data">
+          window.__RELAY_DATA = window.__RELAY_DATA || [];
+          ${this._queryData.map(asRelayDataAppend).join("\n")}
+          Array.prototype.forEach.call(
+            document.getElementsByClassName("__relay_data"),
+            function (element) {
+              element.remove()
+            }
+          );
+          </script>`;
+
+      this._queryData = [];
+    }
+
+    if (this._assetLoaderTags.length > 0) {
+      scriptTags += this._assetLoaderTags.join("");
+
+      this._assetLoaderTags = [];
+    }
+
+    return scriptTags;
+  }
+
+  _write(chunk, encoding, callback) {
+    // This should pick up any new tags that hasn't been previously
+    // written to this stream.
+    let scriptTags = this._generateNewScriptTagsSinceLastCall();
+    if (scriptTags.length > 0) {
+      // Write it before the HTML to ensure that we can start
+      // downloading it as early as possible.
+      this._writable.write(scriptTags);
+    }
+    // Finally write whatever React tried to write.
+    this._writable.write(chunk, encoding, callback);
+  }
+
+  flush() {
+    if (typeof this._writable.flush === 'function') {
+      this._writable.flush();
+    }
+  }
+}

server.mjs

@@ -5,11 +5,8 @@ import { fileURLToPath } from "url";
 import path from "path";
 import fetch from "node-fetch";
 import stream from "stream";
-import {
-  RescriptRelayWritable,
-  writeAssetsIntoStream,
-} from "./streamUtils.mjs";
 import { findGeneratedModule } from "./lookup.mjs";
+import PreloadInsertingStream from "./PreloadInsertingStream.mjs";
 
 global.fetch = fetch;
 
@@ -39,16 +36,9 @@ async function createServer() {
 
       let [start, end] = template.split("<!--ssr-outlet-->");
 
-      let queryDataHolder = { queryData: [] };
-      let preloadAssetHolder = { assets: [] };
-
       let strm = new stream.PassThrough();
 
-      let s = new RescriptRelayWritable(
-        strm,
-        queryDataHolder,
-        preloadAssetHolder
-      );
+      let preloadInsertingStream = new PreloadInsertingStream(strm);
 
       // Pipe everything from our pass through stream into res so it goes to the
       // client.
@@ -57,13 +47,13 @@ async function createServer() {
       // This here is a trade off. It lets us start streaming early, but it also
       // means we'll always return 200 since there's no way we can catch error
       // before starting the stream, as we do it early.
+      // TODO: We can move this into `on*Ready` without any performance penalty. At that point we at least know whether the shell (initial suspense boundaries) was successful.
       res.statusCode = didError ? 500 : 200;
       res.setHeader("Content-type", "text/html");
-      s.write(start);
+      preloadInsertingStream.write(start);
+
+      preloadInsertingStream.on("finish", () => {
 
-      s.on("finish", () => {
-        console.log("[debug] writing end...");
-        strm.write(end);
         strm.end();
       });
 
@@ -81,13 +71,11 @@ async function createServer() {
             // The shell is complete, and React is ready to start streaming.
             // Pipe the results to the intermediate stream.
             console.log("[debug-react-stream] shell completed");
-            writeAssetsIntoStream({
-              queryDataHolder,
-              preloadAssetHolder,
-              writable: s,
-            });
 
-            pipe(s);
+            pipe(preloadInsertingStream);
+
+            console.log("[debug] writing end...");
+            strm.write(end);
           },
           onAllReady() {
             // Write the end of the HTML document when React has streamed
@@ -115,15 +103,15 @@ async function createServer() {
               response,
             })}`
           );
-          queryDataHolder.queryData.push({ id, response, final });
+          preloadInsertingStream.onQuery(id, response, final);
         },
         (id) => {
           console.log(
             `[debug-datalayer] notifying client about started query: ${JSON.stringify(
               id
             )}`
           );
-          queryDataHolder.queryData.push({ id });
+          preloadInsertingStream.onQuery(id)
         },
         // Handle asset preloading. Ideally this should be handled in ReScript
         // code instead, giving that handler the server manifest.
@@ -140,7 +128,7 @@ async function createServer() {
                 const mod = vite.moduleGraph.getModuleById(rescriptModuleLoc);
 
                 if (mod != null) {
-                  preloadAssetHolder.assets.push(
+                  preloadInsertingStream.onAssetPreload(
                     `<script type="module" src="${mod.url}" async></script>`
                   );
                 }