From 95ac5ca4aea7cc707ba22809a77a1f0ca51d21db Mon Sep 17 00:00:00 2001 From: Hiroshi Ogawa Date: Fri, 28 Jun 2024 18:57:14 +0900 Subject: [PATCH 1/3] refactor: simplify ssr flight stream --- packages/react-server/src/entry/browser.tsx | 4 +- packages/react-server/src/entry/server.tsx | 11 +-- .../react-server/src/utils/stream-script.tsx | 86 ++++++++----------- packages/react-server/src/utils/stream.ts | 11 --- 4 files changed, 41 insertions(+), 71 deletions(-) delete mode 100644 packages/react-server/src/utils/stream.ts diff --git a/packages/react-server/src/entry/browser.tsx b/packages/react-server/src/entry/browser.tsx index 69e66c56d..61dd60040 100644 --- a/packages/react-server/src/entry/browser.tsx +++ b/packages/react-server/src/entry/browser.tsx @@ -21,7 +21,7 @@ import { } from "../lib/client/router"; import { $__global } from "../lib/global"; import type { CallServerCallback } from "../lib/types"; -import { readStreamScript } from "../utils/stream-script"; +import { getFlightStreamBrowser } from "../utils/stream-script"; const debug = createDebug("react-server:browser"); @@ -67,7 +67,7 @@ export async function start() { // TODO: needs to await for hydration formState. does it affect startup perf? const initialLayout = await ReactClient.createFromReadableStream( - readStreamScript().pipeThrough(new TextEncoderStream()), + getFlightStreamBrowser(), { callServer }, ); const initialLayoutPromise = Promise.resolve(initialLayout); diff --git a/packages/react-server/src/entry/server.tsx b/packages/react-server/src/entry/server.tsx index 668ce068a..b7b5b126e 100644 --- a/packages/react-server/src/entry/server.tsx +++ b/packages/react-server/src/entry/server.tsx @@ -27,10 +27,9 @@ import { import { $__global } from "../lib/global"; import { ENTRY_REACT_SERVER_WRAPPER, invalidateModule } from "../plugin/utils"; import { escpaeScriptString } from "../utils/escape"; -import { jsonStringifyTransform } from "../utils/stream"; import { createBufferedTransformStream, - injectStreamScript, + injectFlightStream, } from "../utils/stream-script"; import type { ReactServerHandlerStreamResult } from "./react-server"; @@ -204,13 +203,7 @@ export async function renderHtml( .pipeThrough(new TextDecoderStream()) .pipeThrough(createBufferedTransformStream()) .pipeThrough(injectToHead(head)) - .pipeThrough( - injectStreamScript( - stream2 - .pipeThrough(new TextDecoderStream()) - .pipeThrough(jsonStringifyTransform()), - ), - ) + .pipeThrough(injectFlightStream(stream2)) .pipeThrough(new TextEncoderStream()); return new Response(htmlStream, { diff --git a/packages/react-server/src/utils/stream-script.tsx b/packages/react-server/src/utils/stream-script.tsx index f7b8c3f26..31f141411 100644 --- a/packages/react-server/src/utils/stream-script.tsx +++ b/packages/react-server/src/utils/stream-script.tsx @@ -4,66 +4,54 @@ // https://github.com/vercel/next.js/blob/1c5aa7fa09cc5503c621c534fc40065cbd2aefcb/packages/next/src/client/app-index.tsx#L110-L113 // https://github.com/devongovett/rsc-html-stream/ -export function injectStreamScript(stream: ReadableStream) { - const search = ""; +const INIT_SCRIPT = ` +self.__flightStream = new ReadableStream({ + start(controller) { + self.__f_push = (c) => controller.enqueue(c); + self.__f_close = () => controller.close(); + } +}).pipeThrough(new TextEncoderStream()); +`; + +export function injectFlightStream(stream: ReadableStream) { return new TransformStream({ async transform(chunk, controller) { - if (!chunk.includes(search)) { + if (chunk.includes("")) { + controller.enqueue( + chunk.replace( + "", + () => ``, + ), + ); + } else if (chunk.includes("")) { + const i = chunk.indexOf(""); + controller.enqueue(chunk.slice(0, i)); + await stream.pipeThrough(new TextDecoderStream()).pipeTo( + new WritableStream({ + write(chunk) { + controller.enqueue( + ``, + ); + }, + close() { + controller.enqueue(``); + }, + }), + ); + controller.enqueue(chunk.slice(i)); + } else { controller.enqueue(chunk); - return; } - - const [pre, post] = chunk.split(search); - controller.enqueue(pre); - - // TODO: handle cancel? - await stream.pipeTo( - new WritableStream({ - start() { - controller.enqueue(``); - }, - write(chunk) { - // assume chunk is already encoded as javascript code e.g. by - // stream.pipeThrough(jsonStringifyTransform()) - controller.enqueue( - ``, - ); - }, - }), - ); - - controller.enqueue(search + post); }, }); } -export function readStreamScript() { - return new ReadableStream({ - start(controller) { - const chunks: T[] = ((globalThis as any).__stream_chunks ||= []); - - for (const chunk of chunks) { - controller.enqueue(chunk); - } - - chunks.push = function (chunk) { - controller.enqueue(chunk); - return 0; - }; - - if (document.readyState === "loading") { - document.addEventListener("DOMContentLoaded", () => { - controller.close(); - }); - } else { - controller.close(); - } - }, - }); +export function getFlightStreamBrowser(): ReadableStream { + return (self as any).__flightStream; } // it seems buffering is necessary to ensure tag marker (e.g. ``) is not split into multiple chunks. -// Without this, above `injectStreamScript` breaks when receiving two chunks for "...<" and "/body>...". +// Without this, above `injectFlightStream` breaks when receiving two chunks separately for "...<" and "/body>...". // see https://github.com/hi-ogawa/vite-plugins/pull/457 export function createBufferedTransformStream() { let timeout: ReturnType | undefined; diff --git a/packages/react-server/src/utils/stream.ts b/packages/react-server/src/utils/stream.ts deleted file mode 100644 index 6faaba53a..000000000 --- a/packages/react-server/src/utils/stream.ts +++ /dev/null @@ -1,11 +0,0 @@ -function mapTransformStream(f: (v: T) => U) { - return new TransformStream({ - transform(chunk, controller) { - controller.enqueue(f(chunk)); - }, - }); -} - -export function jsonStringifyTransform() { - return mapTransformStream(JSON.stringify); -} From 035d58d9315c1ce0c7f5ff5cbf9e3b0ad59cd104 Mon Sep 17 00:00:00 2001 From: Hiroshi Ogawa Date: Fri, 28 Jun 2024 18:59:05 +0900 Subject: [PATCH 2/3] chore: comment --- packages/react-server/src/entry/server.tsx | 2 -- 1 file changed, 2 deletions(-) diff --git a/packages/react-server/src/entry/server.tsx b/packages/react-server/src/entry/server.tsx index b7b5b126e..7f80cfaf4 100644 --- a/packages/react-server/src/entry/server.tsx +++ b/packages/react-server/src/entry/server.tsx @@ -180,8 +180,6 @@ export async function renderHtml( // render empty as error fallback and // let browser render full CSR instead of hydration // which will replay client error boudnary from RSC error - // TODO: proper two-pass SSR with error route tracking? - // TODO: meta tag system const errorRoot = ( From 9a52e5757c23e18361f6d5f58b543a4b27bda6cb Mon Sep 17 00:00:00 2001 From: Hiroshi Ogawa Date: Fri, 28 Jun 2024 19:04:03 +0900 Subject: [PATCH 3/3] fix: fix injectFlightStream --- packages/react-server/src/utils/stream-script.tsx | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/packages/react-server/src/utils/stream-script.tsx b/packages/react-server/src/utils/stream-script.tsx index 31f141411..73addc279 100644 --- a/packages/react-server/src/utils/stream-script.tsx +++ b/packages/react-server/src/utils/stream-script.tsx @@ -17,13 +17,12 @@ export function injectFlightStream(stream: ReadableStream) { return new TransformStream({ async transform(chunk, controller) { if (chunk.includes("")) { - controller.enqueue( - chunk.replace( - "", - () => ``, - ), + chunk = chunk.replace( + "", + () => ``, ); - } else if (chunk.includes("")) { + } + if (chunk.includes("")) { const i = chunk.indexOf(""); controller.enqueue(chunk.slice(0, i)); await stream.pipeThrough(new TextDecoderStream()).pipeTo(