feat: optimize js loading for streaming ssr

packages/cli/builder/src/index.ts

@@ -31,6 +31,7 @@ export {
 
 export {
   RUNTIME_CHUNK_NAME,
+  RUNTIME_CHUNK_REGEX,
   SERVICE_WORKER_ENVIRONMENT_NAME,
   isHtmlDisabled,
   castArray,

packages/cli/builder/src/plugins/runtimeChunk.ts

@@ -1,5 +1,5 @@
 import type { RsbuildPlugin } from '@rsbuild/core';
-import { RUNTIME_CHUNK_NAME } from '../shared/utils';
+import { RUNTIME_CHUNK_NAME, RUNTIME_CHUNK_REGEX } from '../shared/utils';
 
 export const pluginRuntimeChunk = (
   disableInlineRuntimeChunk?: boolean,
@@ -30,12 +30,8 @@ export const pluginRuntimeChunk = (
         return;
       }
 
-      // RegExp like /bundler-runtime([.].+)?\.js$/
-      // matches bundler-runtime.js and bundler-runtime.123456.js
-      const regexp = new RegExp(`${RUNTIME_CHUNK_NAME}([.].+)?\\.js$`);
-
       if (!config.output.inlineScripts) {
-        config.output.inlineScripts = regexp;
+        config.output.inlineScripts = RUNTIME_CHUNK_REGEX;
       }
     });
   },

packages/cli/builder/src/shared/utils.ts

@@ -3,6 +3,11 @@ import browserslist from 'browserslist';
 
 export const RUNTIME_CHUNK_NAME = 'builder-runtime';
 
+// RegExp like /builder-runtime([.].+)?\.js$/
+export const RUNTIME_CHUNK_REGEX = new RegExp(
+  `${RUNTIME_CHUNK_NAME}([.].+)?\\.js$`,
+);
+
 export const SERVICE_WORKER_ENVIRONMENT_NAME = 'workerSSR';
 
 export const NODE_MODULES_REGEX = /[\\/]node_modules[\\/]/;

packages/runtime/plugin-runtime/src/core/server/stream/createReadableStream.worker.ts

@@ -68,17 +68,45 @@ export const createReadableStreamFromElement: CreateReadableStreamFromElement =
       const stream = new ReadableStream({
         start(controller) {
           const pendingScripts: string[] = [];
+          let isClosed = false;
+
+          const safeEnqueue = (chunk: Uint8Array | unknown) => {
+            if (isClosed) return;
+            try {
+              controller.enqueue(chunk as Uint8Array);
+            } catch {
+              isClosed = true;
+            }
+          };
+
+          const closeController = () => {
+            if (!isClosed) {
+              isClosed = true;
+              try {
+                controller.close();
+              } catch {
+                // Controller already closed
+              }
+            }
+          };
+
+          const flushPendingScripts = () => {
+            for (const s of pendingScripts) {
+              safeEnqueue(encodeForWebStream(s));
+            }
+            pendingScripts.length = 0;
+          };
+
           const enqueueScript = (script: string) => {
             if (shellChunkStatus === ShellChunkStatus.FINISH) {
-              controller.enqueue(encodeForWebStream(script));
+              safeEnqueue(encodeForWebStream(script));
             } else {
               pendingScripts.push(script);
             }
           };
 
           const storageContext = storage.useContext?.();
           const activeDeferreds = storageContext?.activeDeferreds;
-
           /**
            * activeDeferreds is injected into storageContext by @modern-js/runtime.
            * @see packages/toolkit/runtime-utils/src/browser/nestedRoutes.tsx
@@ -89,51 +117,58 @@ export const createReadableStreamFromElement: CreateReadableStreamFromElement =
               : [];
 
           if (entries.length > 0) {
-            enqueueFromEntries(entries, config.nonce, (s: string) =>
-              enqueueScript(s),
-            );
+            enqueueFromEntries(entries, config.nonce, enqueueScript);
           }
 
           async function push() {
-            const { done, value } = await reader.read();
-            if (done) {
-              controller.close();
-              return;
-            }
-            if (shellChunkStatus !== ShellChunkStatus.FINISH) {
-              const chunk = new TextDecoder().decode(value);
-
-              chunkVec.push(chunk);
-
-              let concatedChunk = chunkVec.join('');
-              if (concatedChunk.includes(ESCAPED_SHELL_STREAM_END_MARK)) {
-                concatedChunk = concatedChunk.replace(
-                  ESCAPED_SHELL_STREAM_END_MARK,
-                  '',
-                );
-
-                shellChunkStatus = ShellChunkStatus.FINISH;
-
-                controller.enqueue(
-                  encodeForWebStream(
-                    `${shellBefore}${concatedChunk}${shellAfter}`,
-                  ),
-                );
-                // Flush any pending <script> collected before shell finished
-                if (pendingScripts.length > 0) {
-                  for (const s of pendingScripts) {
-                    controller.enqueue(encodeForWebStream(s));
-                  }
-                  pendingScripts.length = 0;
+            try {
+              const { done, value } = await reader.read();
+              if (done) {
+                closeController();
+                return;
+              }
+
+              if (isClosed) return;
+
+              if (shellChunkStatus !== ShellChunkStatus.FINISH) {
+                chunkVec.push(new TextDecoder().decode(value));
+                const concatedChunk = chunkVec.join('');
+
+                if (concatedChunk.includes(ESCAPED_SHELL_STREAM_END_MARK)) {
+                  shellChunkStatus = ShellChunkStatus.FINISH;
+                  safeEnqueue(
+                    encodeForWebStream(
+                      `${shellBefore}${concatedChunk.replace(
+                        ESCAPED_SHELL_STREAM_END_MARK,
+                        '',
+                      )}${shellAfter}`,
+                    ),
+                  );
+                  flushPendingScripts();
+                }
+              } else {
+                safeEnqueue(value);
+              }
+
+              if (!isClosed) push();
+            } catch (error) {
+              if (!isClosed) {
+                isClosed = true;
+                try {
+                  controller.error(error);
+                } catch {
+                  // Controller already closed
                 }
               }
-            } else {
-              controller.enqueue(value);
             }
-            push();
           }
           push();
         },
+        cancel(reason) {
+          reader.cancel(reason).catch(() => {
+            // Ignore cancellation errors
+          });
+        },
       });
       return stream;
     } catch (e) {

packages/solutions/app-tools/src/builder/shared/builderPlugins/adapterSSR.ts

@@ -29,7 +29,7 @@ export const builderPluginAdapterSSR = (
     api.modifyRsbuildConfig(config => {
       return mergeRsbuildConfig(config, {
         html: {
-          inject: isStreamingSSR(normalizedConfig) ? 'body' : undefined,
+          inject: isStreamingSSR(normalizedConfig) ? 'head' : undefined,
         },
         server: {
           // the http-compression can't handler stream http.

packages/solutions/app-tools/src/builder/shared/bundlerPlugins/HtmlAsyncChunkPlugin.ts

@@ -1,4 +1,5 @@
 import type { Rspack } from '@modern-js/builder';
+import { RUNTIME_CHUNK_REGEX } from '@modern-js/builder';
 
 export class HtmlAsyncChunkPlugin {
   name: string;
@@ -15,19 +16,37 @@ export class HtmlAsyncChunkPlugin {
       const hooks = this.htmlWebpackPlugin.getHooks(compilation as any);
 
       hooks.alterAssetTagGroups.tap(this.name, assets => {
-        const tags = [...assets.headTags, ...assets.bodyTags];
+        const headTags: typeof assets.headTags = [];
+        const bodyTags: typeof assets.bodyTags = [];
 
-        for (const tag of tags) {
+        const processScriptTag = (tag: (typeof assets.headTags)[0]) => {
+          const { attributes } = tag;
+
+          // Convert defer to async
+          if (attributes && attributes.defer === true) {
+            attributes.async = true;
+            delete attributes.defer;
+          }
+
+          const src = attributes?.src as string | undefined;
+          const isRuntimeChunk = src && RUNTIME_CHUNK_REGEX.test(src);
+
+          return isRuntimeChunk ? bodyTags : headTags;
+        };
+
+        for (const tag of [...assets.headTags, ...assets.bodyTags]) {
           if (tag.tagName === 'script') {
-            const { attributes } = tag;
-            if (attributes && attributes.defer === true) {
-              attributes.async = true;
-              delete attributes.defer;
-            }
+            processScriptTag(tag).push(tag);
+          } else {
+            (assets.headTags.includes(tag) ? headTags : bodyTags).push(tag);
           }
         }
 
-        return assets;
+        return {
+          ...assets,
+          headTags,
+          bodyTags,
+        };
       });
     });
   }