refactor(server): refactor compilation middleware

packages/core/src/provider/createCompiler.ts

@@ -11,12 +11,7 @@ import {
 } from '../helpers';
 import { registerDevHook } from '../hooks';
 import { logger } from '../logger';
-import type {
-  DevConfig,
-  InternalContext,
-  Rspack,
-  ServerConfig,
-} from '../types';
+import type { InternalContext, Rspack } from '../types';
 import { type InitConfigsOptions, initConfigs } from './initConfigs';
 
 // keep the last 3 parts of the path to make logs clean
@@ -201,35 +196,3 @@ export async function createCompiler(options: InitConfigsOptions): Promise<{
     rspackConfigs,
   };
 }
-
-export type MiddlewareCallbacks = {
-  onInvalid: () => void;
-  onDone: (stats: any) => void;
-};
-
-export type DevMiddlewareOptions = {
-  /** To ensure HMR works, the devMiddleware need inject the HMR client path into page when HMR enable. */
-  clientPaths?: string[];
-  clientConfig: DevConfig['client'];
-  publicPath?: string;
-
-  /** When liveReload is disabled, the page does not reload. */
-  liveReload?: boolean;
-
-  etag?: 'weak' | 'strong';
-
-  /** The options need by compiler middleware (like webpackMiddleware) */
-  headers?: Record<string, string | string[]>;
-  writeToDisk?:
-    | boolean
-    | ((filename: string, compilationName?: string) => boolean);
-  stats?: boolean;
-
-  /** should trigger when compiler hook called */
-  callbacks: MiddlewareCallbacks;
-
-  /** whether use Server Side Render */
-  serverSideRender?: boolean;
-
-  serverConfig: ServerConfig;
-};

packages/core/src/server/compilerDevMiddleware.ts → packages/core/src/server/compilationManager.ts

@@ -1,7 +1,8 @@
+import fs from 'node:fs';
 import type { IncomingMessage, ServerResponse } from 'node:http';
 import { createRequire } from 'node:module';
-import type { Socket } from 'node:net';
 import { HTML_REGEX } from '../constants';
+import { isMultiCompiler } from '../helpers';
 import { pathnameParse } from '../helpers/path';
 import type {
   EnvironmentContext,
@@ -11,10 +12,9 @@ import type {
   ServerConfig,
 } from '../types';
 import {
-  type DevMiddleware as CustomDevMiddleware,
-  type DevMiddlewareAPI,
-  getDevMiddleware,
-} from './devMiddleware';
+  type CompilationMiddleware,
+  getCompilationMiddleware,
+} from './compilationMiddleware';
 import { stripBase } from './helper';
 import { SocketServer } from './socketServer';
 
@@ -85,8 +85,10 @@ function getClientPaths(devConfig: DevConfig) {
  * 1. setup rsbuild-dev-middleware
  * 2. establish webSocket connect
  */
-export class CompilerDevMiddleware {
-  public middleware!: DevMiddlewareAPI;
+export class CompilationManager {
+  public middleware!: CompilationMiddleware;
+
+  public outputFileSystem: Rspack.OutputFileSystem;
 
   private devConfig: DevConfig;
 
@@ -96,30 +98,27 @@ export class CompilerDevMiddleware {
 
   private publicPaths: string[];
 
-  private socketServer: SocketServer;
+  public socketServer: SocketServer;
 
   constructor({ dev, server, compiler, publicPaths, environments }: Options) {
     this.devConfig = formatDevConfig(dev, environments);
     this.serverConfig = server;
     this.compiler = compiler;
     this.publicPaths = publicPaths;
-
-    // init socket server
+    this.outputFileSystem = fs;
     this.socketServer = new SocketServer(dev);
   }
 
   public async init(): Promise<void> {
-    // start compiling
-    const devMiddleware = await getDevMiddleware(this.compiler);
-    this.middleware = await this.setupDevMiddleware(
-      devMiddleware,
-      this.publicPaths,
-    );
+    await this.setupCompilationMiddleware();
     await this.socketServer.prepare();
-  }
 
-  public upgrade(req: IncomingMessage, sock: Socket, head: any): void {
-    this.socketServer.upgrade(req, sock, head);
+    // Get the latest outputFileSystem from rsbuild-dev-middleware
+    const { compiler } = this;
+    this.outputFileSystem =
+      (isMultiCompiler(compiler)
+        ? compiler.compilers[0].outputFileSystem
+        : compiler.outputFileSystem) || fs;
   }
 
   public async close(): Promise<void> {
@@ -142,22 +141,18 @@ export class CompilerDevMiddleware {
     });
   }
 
-  public sockWrite(
-    type: string,
-    data?: Record<string, any> | string | boolean,
-  ): void {
-    this.socketServer.sockWrite({
-      type,
-      data,
-    });
-  }
+  public readFileSync = (fileName: string): string => {
+    if ('readFileSync' in this.outputFileSystem) {
+      // bundle require needs a synchronous method, although readFileSync is not within the
+      // outputFileSystem type definition, but nodejs fs API implemented.
+      // @ts-expect-error
+      return this.outputFileSystem.readFileSync(fileName, 'utf-8');
+    }
+    return fs.readFileSync(fileName, 'utf-8');
+  };
 
-  private async setupDevMiddleware(
-    devMiddleware: CustomDevMiddleware,
-    publicPaths: string[],
-  ): Promise<DevMiddlewareAPI> {
-    const { devConfig, serverConfig } = this;
-    const { headers, base } = serverConfig;
+  private async setupCompilationMiddleware(): Promise<void> {
+    const { devConfig, serverConfig, publicPaths } = this;
 
     const callbacks = {
       onInvalid: (compilationId?: string, fileName?: string | null) => {
@@ -169,7 +164,6 @@ export class CompilerDevMiddleware {
           });
           return;
         }
-
         this.socketServer.sockWrite({
           type: 'invalid',
           compilationId,
@@ -182,22 +176,14 @@ export class CompilerDevMiddleware {
 
     const clientPaths = getClientPaths(devConfig);
 
-    const middleware = await devMiddleware({
-      headers,
-      publicPath: '/',
-      stats: false,
+    const middleware = await getCompilationMiddleware(this.compiler, {
       callbacks,
-      clientPaths: clientPaths,
-      clientConfig: devConfig.client,
-      liveReload: devConfig.liveReload,
-      writeToDisk: devConfig.writeToDisk,
-      serverSideRender: true,
-      // weak is enough in dev
-      // https://developer.mozilla.org/en-US/docs/Web/HTTP/Conditional_requests#weak_validation
-      etag: 'weak',
+      clientPaths,
+      devConfig,
       serverConfig,
     });
 
+    const { base } = serverConfig;
     const assetPrefixes = publicPaths
       .map(pathnameParse)
       .map((prefix) =>
@@ -230,6 +216,6 @@ export class CompilerDevMiddleware {
 
     // wrap rsbuild-dev-middleware to handle HTML file（without publicPath）
     // maybe we should serve HTML file by sirv
-    return wrapper;
+    this.middleware = wrapper;
   }
 }

packages/core/src/server/devMiddleware.ts → packages/core/src/server/compilationMiddleware.ts

@@ -1,16 +1,10 @@
 import type { IncomingMessage, ServerResponse } from 'node:http';
 import type { Compiler, MultiCompiler } from '@rspack/core';
 import { applyToCompiler } from '../helpers';
-import type { DevMiddlewareOptions } from '../provider/createCompiler';
-import type { DevConfig, NextFunction } from '../types';
+import type { DevConfig, NextFunction, ServerConfig } from '../types';
 import { getCompilationId } from './helper';
 import { getResolvedClientConfig } from './hmrFallback';
 
-type ServerCallbacks = {
-  onInvalid: (compilationId?: string, fileName?: string | null) => void;
-  onDone: (stats: any) => void;
-};
-
 export const isClientCompiler = (compiler: {
   options: {
     target?: Compiler['options']['target'];
@@ -39,9 +33,14 @@ const isNodeCompiler = (compiler: {
   return false;
 };
 
+type ServerCallbacks = {
+  onInvalid: (compilationId?: string, fileName?: string | null) => void;
+  onDone: (stats: any) => void;
+};
+
 export const setupServerHooks = (
   compiler: Compiler,
-  hookCallbacks: ServerCallbacks,
+  { onDone, onInvalid }: ServerCallbacks,
 ): void => {
   // TODO: node SSR HMR is not supported yet
   if (isNodeCompiler(compiler)) {
@@ -51,36 +50,34 @@ export const setupServerHooks = (
   const { compile, invalid, done } = compiler.hooks;
 
   compile.tap('rsbuild-dev-server', () => {
-    hookCallbacks.onInvalid(getCompilationId(compiler));
+    onInvalid(getCompilationId(compiler));
   });
   invalid.tap('rsbuild-dev-server', (fileName) => {
-    hookCallbacks.onInvalid(getCompilationId(compiler), fileName);
+    onInvalid(getCompilationId(compiler), fileName);
   });
-  done.tap('rsbuild-dev-server', hookCallbacks.onDone);
+  done.tap('rsbuild-dev-server', onDone);
 };
 
 function applyHMREntry({
   compiler,
   clientPaths,
-  clientConfig = {},
-  resolvedClientConfig = {},
-  liveReload = true,
+  devConfig,
+  resolvedClientConfig,
 }: {
   compiler: Compiler;
   clientPaths: string[];
-  clientConfig: DevConfig['client'];
+  devConfig: DevConfig;
   resolvedClientConfig: DevConfig['client'];
-  liveReload: DevConfig['liveReload'];
 }) {
   if (!isClientCompiler(compiler)) {
     return;
   }
 
   new compiler.webpack.DefinePlugin({
     RSBUILD_COMPILATION_NAME: JSON.stringify(getCompilationId(compiler)),
-    RSBUILD_CLIENT_CONFIG: JSON.stringify(clientConfig),
+    RSBUILD_CLIENT_CONFIG: JSON.stringify(devConfig.client),
     RSBUILD_RESOLVED_CLIENT_CONFIG: JSON.stringify(resolvedClientConfig),
-    RSBUILD_DEV_LIVE_RELOAD: liveReload,
+    RSBUILD_DEV_LIVE_RELOAD: devConfig.liveReload,
   }).apply(compiler);
 
   for (const clientPath of clientPaths) {
@@ -90,62 +87,71 @@ function applyHMREntry({
   }
 }
 
-export type Middleware = (
+type Middleware = (
   req: IncomingMessage,
   res: ServerResponse,
   next: NextFunction,
 ) => Promise<void>;
 
-export type DevMiddlewareAPI = Middleware & {
+export type CompilationMiddlewareOptions = {
+  /**
+   * To ensure HMR works, the devMiddleware need inject the HMR client path into page when HMR enable.
+   */
+  clientPaths?: string[];
+  /**
+   * Should trigger when compiler hook called
+   */
+  callbacks: ServerCallbacks;
+  devConfig: DevConfig;
+  serverConfig: ServerConfig;
+};
+
+export type CompilationMiddleware = Middleware & {
   close: (callback: (err: Error | null | undefined) => void) => any;
 };
 
 /**
- * The rsbuild/server do nothing about compiler, the devMiddleware need do such things to ensure dev works well:
- * - Call compiler.watch （normally did by rsbuild-dev-middleware）.
- * - Inject the HMR client path into page （the HMR client rsbuild/server already provide）.
- * - Notify server when compiler hooks are triggered.
+ * The CompilationMiddleware handles compiler setup for development:
+ * - Call `compiler.watch` (handled by rsbuild-dev-middleware)
+ * - Inject the HMR client path into page
+ * - Notify server when compiler hooks are triggered
  */
-export type DevMiddleware = (
-  options: DevMiddlewareOptions,
-) => Promise<DevMiddlewareAPI>;
-
-export const getDevMiddleware = async (
-  multiCompiler: Compiler | MultiCompiler,
-): Promise<NonNullable<DevMiddleware>> => {
+export const getCompilationMiddleware = async (
+  compiler: Compiler | MultiCompiler,
+  options: CompilationMiddlewareOptions,
+): Promise<CompilationMiddleware> => {
   const { default: rsbuildDevMiddleware } = await import(
     '../../compiled/rsbuild-dev-middleware/index.js'
   );
-  return async (options) => {
-    const {
-      clientPaths,
-      clientConfig,
-      callbacks,
-      liveReload,
-      serverConfig,
-      ...restOptions
-    } = options;
-    const resolvedClientConfig = await getResolvedClientConfig(
-      clientConfig,
-      serverConfig,
-    );
-
-    const setupCompiler = (compiler: Compiler) => {
-      if (clientPaths) {
-        applyHMREntry({
-          compiler,
-          clientPaths,
-          clientConfig,
-          resolvedClientConfig,
-          liveReload,
-        });
-      }
-      // register hooks for each compilation, update socket stats if recompiled
-      setupServerHooks(compiler, callbacks);
-    };
-
-    applyToCompiler(multiCompiler, setupCompiler);
-
-    return rsbuildDevMiddleware(multiCompiler, restOptions);
+
+  const { clientPaths, callbacks, devConfig, serverConfig } = options;
+  const resolvedClientConfig = await getResolvedClientConfig(
+    devConfig.client,
+    serverConfig,
+  );
+
+  const setupCompiler = (compiler: Compiler) => {
+    if (clientPaths) {
+      applyHMREntry({
+        compiler,
+        clientPaths,
+        devConfig,
+        resolvedClientConfig,
+      });
+    }
+    // register hooks for each compilation, update socket stats if recompiled
+    setupServerHooks(compiler, callbacks);
   };
+
+  applyToCompiler(compiler, setupCompiler);
+
+  return rsbuildDevMiddleware(compiler, {
+    // weak is enough in dev
+    // https://developer.mozilla.org/en-US/docs/Web/HTTP/Conditional_requests#weak_validation
+    etag: 'weak',
+    publicPath: '/',
+    stats: false,
+    serverSideRender: true,
+    writeToDisk: devConfig.writeToDisk,
+  });
 };