rwsdk
Version:
Build fast, server-driven webapps on Cloudflare with SSR, RSC, and realtime
411 lines (410 loc) • 19.7 kB
JavaScript
/**
* A utility to orchestrate and interleave two ReadableStreams (a document shell and an app shell)
* based on a set of markers within their content. This is designed to solve a specific
* race condition in streaming Server-Side Rendering (SSR) with Suspense.
*
* The logic is as follows:
* 1. Stream the document until a start marker is found.
* 2. Switch to the app stream and stream it until an end marker is found. This is the non-suspended shell.
* 3. Switch back to the document stream and stream it until the closing body tag. This sends the client script.
* 4. Switch back to the app stream and stream the remainder (the suspended content).
* 5. Switch back to the document stream and stream the remainder (closing body and html tags).
*
* @param outerHtml The stream for the document shell (`<Document>`).
* @param innerHtml The stream for the application's content.
* @param startMarker The marker in the document to start injecting the app.
* @param endMarker The marker in the app stream that signals the end of the initial, non-suspended render.
*/
function splitStreamOnFirstNonHoistedTag(sourceStream) {
const decoder = new TextDecoder();
const encoder = new TextEncoder();
const nonHoistedTagPattern = /<(?!(?:\/)?(?:title|meta|link|style|base)[\s>\/])(?![!?])/i;
let sourceReader;
let appBodyController = null;
let buffer = "";
let hoistedTagsDone = false;
const hoistedTagsStream = new ReadableStream({
start(controller) {
sourceReader = sourceStream.getReader();
const pump = async () => {
try {
if (hoistedTagsDone) {
controller.close();
return;
}
const { done, value } = await sourceReader.read();
if (done) {
if (buffer) {
const match = buffer.match(nonHoistedTagPattern);
if (match && typeof match.index === "number") {
const hoistedPart = buffer.slice(0, match.index);
controller.enqueue(encoder.encode(hoistedPart));
}
else {
controller.enqueue(encoder.encode(buffer));
}
}
controller.close();
hoistedTagsDone = true;
if (appBodyController) {
appBodyController.close();
}
return;
}
buffer += decoder.decode(value, { stream: true });
const match = buffer.match(nonHoistedTagPattern);
if (match && typeof match.index === "number") {
const hoistedPart = buffer.slice(0, match.index);
const appPart = buffer.slice(match.index);
buffer = "";
controller.enqueue(encoder.encode(hoistedPart));
controller.close();
hoistedTagsDone = true;
if (appBodyController) {
if (appPart) {
appBodyController.enqueue(encoder.encode(appPart));
}
while (true) {
const { done, value } = await sourceReader.read();
if (done) {
appBodyController.close();
return;
}
appBodyController.enqueue(value);
}
}
}
else {
const flushIndex = buffer.lastIndexOf("\n");
if (flushIndex !== -1) {
controller.enqueue(encoder.encode(buffer.slice(0, flushIndex + 1)));
buffer = buffer.slice(flushIndex + 1);
}
await pump();
}
}
catch (e) {
controller.error(e);
if (appBodyController) {
appBodyController.error(e);
}
}
};
pump().catch((e) => {
controller.error(e);
if (appBodyController) {
appBodyController.error(e);
}
});
},
});
const appBodyStream = new ReadableStream({
start(controller) {
appBodyController = controller;
},
});
return [hoistedTagsStream, appBodyStream];
}
/**
* A utility that orchestrates and interleaves three ReadableStreams to produce a
* single, valid HTML response stream. It uses two special markers:
*
* - `startMarker`: Placed in the `outerHtml` stream (the document shell) to
* designate where the application's content should be injected.
* - `endMarker`: Injected into the `innerHtml` stream's RSC payload to signal
* the end of the initial, non-suspended render. This marker is needed for
* non-blocking hydration, as it allows the stitching process to send the
* client `<script>` tags before all suspended content has resolved.
*
* It manages three main stream readers:
*
* - `hoistedTagsReader`: Reads from the `hoistedTagsStream`, which contains only
* the hoisted meta tags (e.g., `<title>`, `<meta>`).
* - `outerReader`: Reads from the `outerHtml` stream, which is the server-rendered
* document shell (containing `<html>`, `<head>`, etc.).
* - `innerReader`: Reads from the `appBodyStream`, which contains the main
* application content, stripped of its hoisted tags.
*
* The function proceeds through a multi-phase state machine, managed by the
* `pump` function, to correctly interleave these streams.
*
* The state machine moves through the following phases:
*
* 1. `read-hoisted`:
* - **Goal:** Buffer all hoisted tags from the `hoistedTagsStream`.
* - **Action:** Reads from `hoistedTagsReader` and appends all content into
* the `hoistedTagsBuffer`. Does not enqueue anything yet.
* - **Transition:** Moves to `outer-head` when the stream is exhausted.
*
* 2. `outer-head`:
* - **Goal:** Stream the document up to the closing `</head>` tag, inject the
* hoisted tags, and then continue until the app `startMarker`.
* - **Action:** Reads from `outerReader`. When it finds `</head>`, it enqueues
* the content before it, then enqueues the `hoistedTagsBuffer`, and finally
* enqueues the `</head>` tag itself. It then continues reading from
* `outerReader` until it finds the `startMarker`.
* - **Transition:** Moves to `inner-shell` after finding and discarding the
* `startMarker`.
*
* 3. `inner-shell`:
* - **Goal:** Stream the initial, non-suspended part of the application.
* - **Action:** Switches to `innerReader`. It enqueues chunks until it finds
* the `endMarker`. Any content after the marker is stored in
* `innerSuspendedRemains`.
* - **Transition:** Moves to `outer-tail` after finding the `endMarker`.
*
* 4. `outer-tail`:
* - **Goal:** Stream the rest of the document's `<body>`, including client
* `<script>` tags.
* - **Action:** Switches back to `outerReader` and enqueues chunks until it
* finds the `</body>` tag.
* - **Transition:** Moves to `inner-suspended` after finding `</body>`.
*
* 5. `inner-suspended`:
* - **Goal:** Stream any suspended content from the React app.
* - **Action:** First enqueues any content from `innerSuspendedRemains`, then
* continues reading from `innerReader` until the stream is exhausted.
* - **Transition:** Moves to `outer-end` when the stream is exhausted.
*
* 6. `outer-end`:
* - **Goal:** Finish the document.
* - **Action:** Switches back to `outerReader` for the last time to send the
* closing `</body>` and `</html>` tags.
*/
export function stitchDocumentAndAppStreams(outerHtml, innerHtml, startMarker, endMarker) {
const [hoistedTagsStream, appBodyStream] = splitStreamOnFirstNonHoistedTag(innerHtml);
const decoder = new TextDecoder();
const encoder = new TextEncoder();
let outerReader;
let innerReader;
let hoistedTagsReader;
let buffer = "";
let outerBufferRemains = "";
let innerSuspendedRemains = "";
let hoistedTagsBuffer = "";
let hoistedTagsReady = false;
let phase = "read-hoisted";
const pump = async (controller) => {
const enqueue = (text) => {
if (text) {
controller.enqueue(encoder.encode(text));
}
};
const flush = () => {
const flushIndex = buffer.lastIndexOf("\n");
if (flushIndex !== -1) {
enqueue(buffer.slice(0, flushIndex + 1));
buffer = buffer.slice(flushIndex + 1);
}
};
try {
if (phase === "read-hoisted") {
// Continuously read from the hoisted tags stream and buffer the
// content. Once the stream is finished, transition to the next phase.
const { done, value } = await hoistedTagsReader.read();
// When the stream is done, we're ready to process the document head.
if (done) {
hoistedTagsReady = true;
phase = "outer-head";
}
else {
// Otherwise, keep appending to the buffer.
hoistedTagsBuffer += decoder.decode(value, { stream: true });
}
}
else if (phase === "outer-head") {
// Read from the document stream. Search for the closing `</head>` tag
// to inject the buffered hoisted tags. Then, search for the
// `startMarker` to know when to start injecting the app shell. Once
// the marker is found, transition to the next phase.
const { done, value } = await outerReader.read();
// Handle the case where the document stream ends.
if (done) {
// If there's content left in the buffer, process it for markers.
if (buffer) {
const headCloseIndex = buffer.indexOf("</head>");
if (headCloseIndex !== -1 &&
hoistedTagsReady &&
hoistedTagsBuffer) {
enqueue(buffer.slice(0, headCloseIndex));
enqueue(hoistedTagsBuffer);
hoistedTagsBuffer = "";
enqueue("</head>");
buffer = buffer.slice(headCloseIndex + "</head>".length);
}
const markerIndex = buffer.indexOf(startMarker);
if (markerIndex !== -1) {
enqueue(buffer.slice(0, markerIndex));
outerBufferRemains = buffer.slice(markerIndex + startMarker.length);
}
else {
enqueue(buffer);
}
buffer = "";
}
else if (hoistedTagsReady && hoistedTagsBuffer) {
enqueue(hoistedTagsBuffer);
hoistedTagsBuffer = "";
}
// Even if the stream ends, we must proceed to the app shell phase.
phase = "inner-shell";
}
else {
// As chunks arrive, append them to the buffer.
buffer += decoder.decode(value, { stream: true });
// Search for the closing head tag to inject hoisted tags.
const headCloseIndex = buffer.indexOf("</head>");
if (headCloseIndex !== -1 && hoistedTagsReady && hoistedTagsBuffer) {
enqueue(buffer.slice(0, headCloseIndex));
enqueue(hoistedTagsBuffer);
hoistedTagsBuffer = "";
enqueue("</head>");
buffer = buffer.slice(headCloseIndex + "</head>".length);
}
// Search for the start marker to switch to the app stream.
const markerIndex = buffer.indexOf(startMarker);
if (markerIndex !== -1) {
enqueue(buffer.slice(0, markerIndex));
outerBufferRemains = buffer.slice(markerIndex + startMarker.length);
buffer = "";
phase = "inner-shell";
}
else {
// If no marker is found yet, flush the buffer up to the last
// newline to keep the stream flowing.
flush();
}
}
}
else if (phase === "inner-shell") {
// Now read from the app stream. We send the initial part of the app
// content until we find the `endMarker`. This marker tells us that the
// non-suspended part of the app is rendered. Any content after this
// marker is considered suspended and is buffered. Then, transition.
const { done, value } = await innerReader.read();
// Handle the case where the app stream ends.
if (done) {
if (buffer)
enqueue(buffer);
phase = "outer-tail";
}
else {
// As chunks arrive, append them to the buffer.
buffer += decoder.decode(value, { stream: true });
const markerIndex = buffer.indexOf(endMarker);
// If the end marker is found, enqueue content up to the marker,
// buffer the rest, and switch to the document tail phase.
if (markerIndex !== -1) {
const endOfMarkerIndex = markerIndex + endMarker.length;
enqueue(buffer.slice(0, endOfMarkerIndex));
innerSuspendedRemains = buffer.slice(endOfMarkerIndex);
buffer = "";
phase = "outer-tail";
}
else {
// If no marker is found yet, flush the buffer.
flush();
}
}
}
else if (phase === "outer-tail") {
// Switch back to the document stream. The goal is to send the rest of
// the document's body, which critically includes the client-side
// `<script>` tags for hydration. We stream until we find the closing
// `</body>` tag and then transition.
// First, process any leftover buffer from the `outer-head` phase.
if (outerBufferRemains) {
buffer = outerBufferRemains;
outerBufferRemains = "";
}
const { done, value } = await outerReader.read();
// Handle the case where the document stream ends.
if (done) {
if (buffer) {
// Search the remaining buffer for the closing body tag.
const markerIndex = buffer.indexOf("</body>");
if (markerIndex !== -1) {
enqueue(buffer.slice(0, markerIndex));
buffer = buffer.slice(markerIndex);
}
else {
enqueue(buffer);
buffer = "";
}
}
// Proceed to the suspended content phase.
phase = "inner-suspended";
}
else {
// As chunks arrive, append them to the buffer.
buffer += decoder.decode(value, { stream: true });
// Search for the closing body tag to switch to suspended content.
const markerIndex = buffer.indexOf("</body>");
if (markerIndex !== -1) {
enqueue(buffer.slice(0, markerIndex));
buffer = buffer.slice(markerIndex);
phase = "inner-suspended";
}
else {
// If no marker is found yet, flush the buffer.
flush();
}
}
}
else if (phase === "inner-suspended") {
// Switch back to the app stream. First, send any buffered suspended
// content from the `inner-shell` phase. Then, stream the rest of the
// app content until it's finished. This is all the content that was
// behind a `<Suspense>` boundary.
// First, send any buffered suspended content from the `inner-shell` phase.
if (innerSuspendedRemains) {
enqueue(innerSuspendedRemains);
innerSuspendedRemains = "";
}
const { done, value } = await innerReader.read();
// When the app stream is done, transition to the final phase.
if (done) {
phase = "outer-end";
}
else {
// Otherwise, pass through the remaining app content directly.
controller.enqueue(value);
}
}
else if (phase === "outer-end") {
// Finally, switch back to the document stream one last time to send
// the closing `</body>` and `</html>` tags and finish the response.
// First, send any leftover buffer from the `outer-tail` phase.
if (buffer) {
enqueue(buffer);
buffer = "";
}
const { done, value } = await outerReader.read();
// When the document stream is done, we're finished.
if (done) {
controller.close();
return;
}
// Otherwise, pass through the final document content.
controller.enqueue(value);
}
await pump(controller);
}
catch (e) {
controller.error(e);
}
};
return new ReadableStream({
start(controller) {
outerReader = outerHtml.getReader();
innerReader = appBodyStream.getReader();
hoistedTagsReader = hoistedTagsStream.getReader();
pump(controller).catch((e) => controller.error(e));
},
cancel(reason) {
outerReader?.cancel(reason);
innerReader?.cancel(reason);
hoistedTagsReader?.cancel(reason);
},
});
}