feat: add a new server hook as afterStreamingRender (#5101)

.changeset/tiny-needles-dance.md

@@ -0,0 +1,8 @@
+---
+'@modern-js/prod-server': minor
+'@modern-js/types': minor
+'@modern-js/server-core': minor
+---
+
+feat: SSR server support afterStreamingRender
+feat: SSR 服务端支持 afterStreamingRender

packages/server/core/src/plugin.ts

@@ -23,6 +23,7 @@ import type {
   ServerRoute,
   HttpMethodDecider,
   ServerInitHookContext,
+  AfterStreamingRenderContext,
 } from '@modern-js/types';
 
 import type { BffUserConfig, ServerOptions, UserConfig } from './types/config';
@@ -152,6 +153,11 @@ const beforeRender = createAsyncPipeline<
 
 const afterRender = createAsyncPipeline<AfterRenderContext, any>();
 
+const afterStreamingRender = createAsyncPipeline<
+  AfterStreamingRenderContext,
+  string
+>();
+
 const beforeSend = createAsyncPipeline<ModernServerContext, RequestResult>();
 
 const afterSend = createParallelWorkflow<{
@@ -212,6 +218,7 @@ const serverHooks = {
   renderToString,
   beforeRender,
   afterRender,
+  afterStreamingRender,
   beforeSend,
   afterSend,
   reset,

packages/server/prod-server/src/libs/hook-api/afterRenderForStream.ts

@@ -0,0 +1,12 @@
+import { Transform } from 'stream';
+import { MaybeAsync } from '@modern-js/plugin';
+
+export const afterRenderInjectableStream = (
+  fn: (content: string) => MaybeAsync<string>,
+) =>
+  new Transform({
+    async write(chunk, _, callback) {
+      this.push(await fn(chunk.toString()));
+      callback();
+    },
+  });

packages/server/prod-server/src/libs/hook-api/index.ts

@@ -5,6 +5,7 @@ import type {
   AfterRenderContext,
   MiddlewareContext,
   ServerRoute,
+  AfterStreamingRenderContext,
 } from '@modern-js/types';
 import { RouteAPI } from './route';
 import { TemplateAPI } from './template';
@@ -46,6 +47,20 @@ export const createAfterRenderContext = (
   };
 };
 
+export const createAfterStreamingRenderContext = (
+  context: ModernServerContext,
+  route: Partial<ServerRoute>,
+): ((chunk: string) => AfterStreamingRenderContext) => {
+  const baseContext = base(context);
+  return (chunk: string) => {
+    return {
+      ...baseContext,
+      route,
+      chunk,
+    };
+  };
+};
+
 export const createMiddlewareContext = (
   context: ModernServerContext,
 ): MiddlewareContext => {

packages/server/prod-server/src/libs/render/index.ts

@@ -10,6 +10,7 @@ import { shouldFlushServerHeader } from '../preload/shouldFlushServerHeader';
 import { handleDirectory } from './static';
 import * as ssr from './ssr';
 import { injectServerData } from './utils';
+import { SSRRenderOptions } from './ssr';
 
 export type RenderHandler = (options: {
   ctx: ModernServerContext;
@@ -100,11 +101,9 @@ export const createRenderHandler: CreateRenderHandler = ({
             },
           });
         }
-        const ssrRenderOptions = {
+        const ssrRenderOptions: SSRRenderOptions = {
           distDir,
-          entryName: route.entryName,
-          urlPath: route.urlPath,
-          bundle: route.bundle,
+          route,
           template: content.toString(),
           staticGenerate,
           nonce,

packages/server/prod-server/src/libs/render/ssr.ts

@@ -8,35 +8,37 @@ import {
 } from '@modern-js/utils';
 import type { ModernServerContext } from '@modern-js/types';
 import { RenderResult, ServerHookRunner } from '../../type';
+import { createAfterStreamingRenderContext } from '../hook-api';
+import { afterRenderInjectableStream } from '../hook-api/afterRenderForStream';
+import type { ModernRoute } from '../route';
 import { RenderFunction, SSRServerContext } from './type';
 import { createLogger, createMetrics } from './measure';
 import { injectServerDataStream, injectServerData } from './utils';
 import { ssrCache } from './ssrCache';
 
+export type SSRRenderOptions = {
+  distDir: string;
+  template: string;
+  route: ModernRoute;
+  staticGenerate: boolean;
+  enableUnsafeCtx?: boolean;
+  nonce?: string;
+};
+
 export const render = async (
   ctx: ModernServerContext,
-  renderOptions: {
-    distDir: string;
-    bundle: string;
-    urlPath: string;
-    template: string;
-    entryName: string;
-    staticGenerate: boolean;
-    enableUnsafeCtx?: boolean;
-    nonce?: string;
-  },
+  renderOptions: SSRRenderOptions,
   runner: ServerHookRunner,
 ): Promise<RenderResult> => {
   const {
-    urlPath,
-    bundle,
     distDir,
+    route,
     template,
-    entryName,
     staticGenerate,
     enableUnsafeCtx = false,
     nonce,
   } = renderOptions;
+  const { urlPath, bundle, entryName } = route;
   const bundleJS = path.join(distDir, bundle);
   const loadableUri = path.join(distDir, LOADABLE_STATS_FILE);
   const loadableStats = fs.existsSync(loadableUri) ? require(loadableUri) : '';
@@ -106,9 +108,23 @@ export const render = async (
       contentType: mime.contentType('html') as string,
     };
   } else {
+    let contentStream = injectServerDataStream(content, ctx);
+    const afterStreamingRenderContext = createAfterStreamingRenderContext(
+      ctx,
+      route,
+    );
+    contentStream = contentStream.pipe(
+      afterRenderInjectableStream((chunk: string) => {
+        const context = afterStreamingRenderContext(chunk);
+        return runner.afterStreamingRender(context, {
+          onLast: ({ chunk }) => chunk,
+        });
+      }),
+    );
+
     return {
       content: '',
-      contentStream: injectServerDataStream(content, ctx),
+      contentStream,
       contentType: mime.contentType('html') as string,
     };
   }

packages/server/prod-server/tests/hook.test.ts

@@ -6,6 +6,7 @@ import {
   createAfterMatchContext,
   createAfterRenderContext,
   createMiddlewareContext,
+  createAfterStreamingRenderContext,
 } from '../src/libs/hook-api';
 import { createContext } from '../src/libs/context';
 import { createDoc } from './helper';
@@ -144,4 +145,25 @@ describe('test hook api', () => {
     response.locals.foo = 'bar';
     expect((locals as any).foo).toBe('bar');
   });
+
+  test('should after streaming render context work correctly', () => {
+    const content = createDoc();
+    const req = httpMocks.createRequest({
+      url: '/',
+      headers: {
+        host: 'modernjs.com',
+      },
+      eventEmitter: Readable,
+      method: 'GET',
+    });
+    const res = httpMocks.createResponse({ eventEmitter: EventEmitter });
+    const afterStreamingRenderContext = createAfterStreamingRenderContext(
+      createContext(req, res),
+      {},
+    );
+
+    const context = afterStreamingRenderContext(content);
+
+    expect(context.chunk).toMatch(content);
+  });
 });

packages/server/prod-server/tests/render.test.ts

@@ -18,11 +18,13 @@ describe('test render function', () => {
         req: {},
       } as any,
       {
-        urlPath: '/foo',
-        bundle: 'bundle.js',
+        route: {
+          urlPath: '/foo',
+          bundle: 'bundle.js',
+          entryName: 'foo',
+        },
         distDir: path.join(__dirname, 'fixtures', 'ssr'),
         template: 'tpl.html',
-        entryName: 'foo',
         staticGenerate: false,
       } as any,
       {

packages/server/server/src/server/workerSSRRender.ts

@@ -3,19 +3,20 @@ import { ServerHookRunner } from '@modern-js/prod-server';
 import axios from 'axios';
 import { mime } from '@modern-js/utils';
 import { ModernServerContext } from '@modern-js/types/server';
+import { ModernRoute } from '@modern-js/prod-server/src/libs/route';
 
 const PORT = 9230;
 
 export async function workerSSRRender(
   ctx: ModernServerContext,
   renderOptions: {
-    urlPath: string;
+    route: ModernRoute;
     [props: string]: any;
   },
   _runner: ServerHookRunner,
 ) {
   const { headers, params } = ctx;
-  const { urlPath } = renderOptions;
+  const { urlPath } = renderOptions.route;
   const url = `http://0.0.0.0:${PORT}/${urlPath}`;
   const resposne = await axios.get(url, {
     timeout: 5000,

packages/toolkit/types/server/hook.d.ts

@@ -68,6 +68,16 @@ export type AfterRenderContext = HookContext & {
   };
 };
 
+export type AfterStreamingRenderContext = HookContext & {
+  route?: Partial<
+    Pick<
+      ServerRoute,
+      'entryName' | 'bundle' | 'isSPA' | 'isSSR' | 'urlPath' | 'entryPath'
+    >
+  >;
+  chunk: string;
+};
+
 export type MiddlewareContext<T extends 'worker' | 'node' = 'node'> =
   HookContext & {
     reporter?: Reporter;