fix(plugin-runtime): streaming ssr split chunk so that can't match SHELL_STREAM_END_MARK (#3719)

* fix(plugin-runtime): streaming ssr split chunk so that cant match SHELL_STREAM_END_MARK

* fix: update test && handler worker streaming ssr

Author: GiveMe-A-Name

.changeset/tame-sloths-look.md

@@ -0,0 +1,6 @@
+---
+'@modern-js/runtime': patch
+---
+
+fix(plugin-runtime): streaming ssr split chunk so that can't match `SHELL_STREAM_END_MARK`
+fix(plugin-runtime): streaming ssr chunk 进行分割导致无法匹配 `SHELL_STREAM_END_MARK`

packages/runtime/plugin-runtime/src/router/runtime/DeferredDataScripts.node.tsx

@@ -181,7 +181,7 @@ const DeferredDataScript = ({
             <ErrorDeferredDataScript routeId={routeId} dataKey={dataKey} />
           }
         >
-          {data => (
+          {(data: any) => (
             <script
               async
               suppressHydrationWarning

packages/runtime/plugin-runtime/src/ssr/serverRender/renderToStream/renderToPipe.ts

@@ -7,19 +7,19 @@ import { getTemplates } from './template';
 export type Pipe<T extends Writable> = (output: T) => Promise<T | string>;
 
 enum ShellChunkStatus {
-  IDLE = 0,
-  START = 1,
-  FINIESH = 2,
+  START = 0,
+  FINIESH = 1,
 }
 
 function renderToPipe(
   rootElement: React.ReactElement,
   context: RuntimeContext,
   options?: RenderToPipeableStreamOptions,
 ) {
-  let shellChunkStatus = ShellChunkStatus.IDLE;
+  let shellChunkStatus = ShellChunkStatus.START;
 
   const { ssrContext } = context;
+  const chunkVec: string[] = [];
   const forUserPipe: Pipe<Writable> = stream => {
     return new Promise(resolve => {
       let renderToPipeableStream;
@@ -40,28 +40,24 @@ function renderToPipe(
             transform(chunk, _encoding, callback) {
               try {
                 if (shellChunkStatus !== ShellChunkStatus.FINIESH) {
-                  let concatedChunk = chunk.toString();
-                  if (shellChunkStatus === ShellChunkStatus.IDLE) {
-                    concatedChunk = `${shellBefore}${concatedChunk}`;
-                    shellChunkStatus = ShellChunkStatus.START;
-                  }
+                  chunkVec.push(chunk.toString());
+
                   /**
                    * The shell content of App may be splitted by multiple chunks to transform,
                    * when any node value's size is larger than the React limitation, refer to:
                    * https://github.com/facebook/react/blob/v18.2.0/packages/react-server/src/ReactServerStreamConfigNode.js#L53.
                    * So we use the `SHELL_STREAM_END_MARK` to mark the shell content' tail.
                    */
-                  if (
-                    shellChunkStatus === ShellChunkStatus.START &&
-                    concatedChunk.endsWith(ESCAPED_SHELL_STREAM_END_MARK)
-                  ) {
+                  let concatedChunk = chunkVec.join('');
+                  if (concatedChunk.endsWith(ESCAPED_SHELL_STREAM_END_MARK)) {
                     concatedChunk = concatedChunk.replace(
                       ESCAPED_SHELL_STREAM_END_MARK,
-                      shellAfter,
+                      '',
                     );
+
                     shellChunkStatus = ShellChunkStatus.FINIESH;
+                    this.push(`${shellBefore}${concatedChunk}${shellAfter}`);
                   }
-                  this.push(concatedChunk);
                 } else {
                   this.push(chunk);
                 }

packages/runtime/plugin-runtime/src/ssr/serverRender/renderToStream/renderToPipe.worker.ts

@@ -10,17 +10,17 @@ import { getTemplates } from './template';
 export type Pipe<T extends Writable> = (output: T) => Promise<T | string>;
 
 enum ShellChunkStatus {
-  IDLE = 0,
-  START = 1,
-  FINIESH = 2,
+  START = 0,
+  FINIESH = 1,
 }
 
 function renderToPipe(
   rootElement: React.ReactElement,
   context: RuntimeContext,
   options?: RenderToReadableStreamOptions,
 ) {
-  let shellChunkStatus = ShellChunkStatus.IDLE;
+  let shellChunkStatus = ShellChunkStatus.START;
+  const chunkVec: string[] = [];
 
   const { ssrContext } = context;
   const forUserPipe = async () => {
@@ -56,20 +56,24 @@ function renderToPipe(
               return;
             }
             if (shellChunkStatus !== ShellChunkStatus.FINIESH) {
-              let concatedChunk = new TextDecoder().decode(value);
-              if (shellChunkStatus === ShellChunkStatus.IDLE) {
-                concatedChunk = `${shellBefore}${concatedChunk}`;
-                shellChunkStatus = ShellChunkStatus.START;
-              }
-              if (
-                shellChunkStatus === ShellChunkStatus.START &&
-                concatedChunk.endsWith(ESCAPED_SHELL_STREAM_END_MARK)
-              ) {
+              const chunk = new TextDecoder().decode(value);
+
+              chunkVec.push(chunk);
+
+              let concatedChunk = chunkVec.join('');
+              if (concatedChunk.endsWith(ESCAPED_SHELL_STREAM_END_MARK)) {
                 concatedChunk = concatedChunk.replace(
                   ESCAPED_SHELL_STREAM_END_MARK,
-                  shellAfter!,
+                  '',
                 );
+
                 shellChunkStatus = ShellChunkStatus.FINIESH;
+
+                controller.enqueue(
+                  encodeForWebStream(
+                    `${shellBefore}${concatedChunk}${shellAfter}`,
+                  ),
+                );
               }
               controller.enqueue(encodeForWebStream(concatedChunk));
             } else {

packages/runtime/plugin-runtime/src/ssr/serverRender/renderToStream/type.ts

@@ -1,4 +1,4 @@
 export type InjectTemplate = {
-  shellBefore?: string;
-  shellAfter?: string;
+  shellBefore: string;
+  shellAfter: string;
 };

packages/runtime/plugin-runtime/tests/ssr/fixtures/streaming-ssr/App.tsx

@@ -4,12 +4,15 @@ const Home = lazy(() => import('./Home'));
 
 const App: React.FunctionComponent = () => {
   return (
-    <div>
-      <div>App Layout</div>
-      <Suspense fallback={<div>loading home...</div>}>
-        <Home />
-      </Suspense>
-    </div>
+    <>
+      <div>
+        <div>App Layout</div>
+        <Suspense fallback={<div>loading home...</div>}>
+          <Home />
+        </Suspense>
+      </div>
+      {'<!--<?- SHELL_STREAM_END ?>-->'}
+    </>
   );
 };