feat: make environment API result cacheable (#4126)

Co-authored-by: neverland <[email protected]>

e2e/cases/server/ssr/index.test.ts

@@ -1,19 +1,28 @@
-import { dev, rspackOnlyTest } from '@e2e/helper';
+import { dev, proxyConsole, rspackOnlyTest } from '@e2e/helper';
 import { expect } from '@playwright/test';
 
 rspackOnlyTest('support SSR', async ({ page }) => {
+  const { logs, restore } = proxyConsole('log');
+
   const rsbuild = await dev({
     cwd: __dirname,
     rsbuildConfig: {},
   });
 
-  const url1 = new URL(`http://localhost:${rsbuild.port}`);
+  const url = new URL(`http://localhost:${rsbuild.port}`);
 
-  const res = await page.goto(url1.href);
+  const res = await page.goto(url.href);
 
   expect(await res?.text()).toMatch(/Rsbuild with React/);
 
+  await page.goto(url.href);
+
+  // bundle result should cacheable and only load once.
+  expect(logs.filter((log) => log.includes('load SSR')).length).toBe(1);
+
   await rsbuild.close();
+
+  restore();
 });
 
 rspackOnlyTest('support SSR with esm target', async ({ page }) => {

e2e/cases/server/ssr/src/index.server.tsx

@@ -3,6 +3,8 @@ import ReactDOMServer from 'react-dom/server';
 import App from './App';
 import { assert } from './assert.server';
 
+console.log('load SSR');
+
 // test dynamic import
 import('./test');
 

packages/core/src/server/devServer.ts

@@ -16,7 +16,11 @@ import type {
 } from '../types';
 import { isCliShortcutsEnabled, setupCliShortcuts } from './cliShortcuts';
 import { CompilerDevMiddleware } from './compilerDevMiddleware';
-import { getTransformedHtml, loadBundle } from './environment';
+import {
+  createCacheableFunction,
+  getTransformedHtml,
+  loadBundle,
+} from './environment';
 import {
   type RsbuildDevMiddlewareOptions,
   getMiddlewares,
@@ -267,6 +271,11 @@ export async function createDevServer<
     return fs.readFileSync(fileName, 'utf-8');
   };
 
+  const cacheableLoadBundle = createCacheableFunction(loadBundle);
+  const cacheableTransformedHtml = createCacheableFunction<string>(
+    (_stats, entryName, utils) => getTransformedHtml(entryName, utils),
+  );
+
   const environmentAPI = Object.fromEntries(
     Object.entries(options.context.environments).map(([name, environment]) => {
       return [
@@ -283,17 +292,25 @@ export async function createDevServer<
           },
           loadBundle: async <T>(entryName: string) => {
             await waitFirstCompileDone;
-            return loadBundle<T>(lastStats[environment.index], entryName, {
-              readFileSync,
-              environment,
-            });
+            return cacheableLoadBundle(
+              lastStats[environment.index],
+              entryName,
+              {
+                readFileSync,
+                environment,
+              },
+            ) as T;
           },
           getTransformedHtml: async (entryName: string) => {
             await waitFirstCompileDone;
-            return getTransformedHtml(entryName, {
-              readFileSync,
-              environment,
-            });
+            return cacheableTransformedHtml(
+              lastStats[environment.index],
+              entryName,
+              {
+                readFileSync,
+                environment,
+              },
+            );
           },
         },
       ];

packages/core/src/server/environment.ts

@@ -78,3 +78,37 @@ export const getTransformedHtml = async (
 
   return fileContent;
 };
+
+export const createCacheableFunction = <T>(
+  getter: (
+    stats: Rspack.Stats,
+    entryName: string,
+    utils: ServerUtils,
+  ) => Promise<T>,
+) => {
+  const cache = new WeakMap<
+    Rspack.Stats,
+    {
+      [entryName: string]: T;
+    }
+  >();
+
+  return async (
+    stats: Rspack.Stats,
+    entryName: string,
+    utils: ServerUtils,
+  ): Promise<T> => {
+    const cachedEntries = cache.get(stats);
+    if (cachedEntries?.[entryName]) {
+      return cachedEntries[entryName];
+    }
+
+    const res = await getter(stats, entryName, utils);
+
+    cache.set(stats, {
+      ...(cachedEntries || {}),
+      [entryName]: res,
+    });
+    return res;
+  };
+};