Refactor build context for better plugin state sharing

packages/next/src/build/build-context.ts

@@ -5,24 +5,49 @@ import type { __ApiPreviewProps } from '../server/api-utils'
 import type { NextConfigComplete } from '../server/config-shared'
 import type { Span } from '../trace'
 import type getBaseWebpackConfig from './webpack-config'
-import { PagesManifest } from './webpack/plugins/pages-manifest-plugin'
+import type { PagesManifest } from './webpack/plugins/pages-manifest-plugin'
 import type { TelemetryPlugin } from './webpack/plugins/telemetry-plugin'
 
+// A layer for storing data that is used by plugins to communicate with each
+// other between different steps of the build process. This is only internal
+// to Next.js and will not be a part of the final build output.
+// These states don't need to be deeply merged.
+let pluginState: Record<string, any> = {}
+export function resumePluginState(resumedState?: Record<string, any>) {
+  Object.assign(pluginState, resumedState)
+}
+
+// This method gives you the plugin state with typed and mutable value fields
+// behind a proxy so we can lazily initialize the values **after** resuming the
+// plugin state.
+export function getProxiedPluginState<State extends Record<string, any>>(
+  initialState: State
+) {
+  return new Proxy(pluginState, {
+    get(target, key: string) {
+      if (typeof target[key] === 'undefined') {
+        return (target[key] = initialState[key])
+      }
+      return target[key]
+    },
+    set(target, key: string, value) {
+      target[key] = value
+      return true
+    },
+  }) as State
+}
+
+export function getPluginState() {
+  return pluginState
+}
+
 // a global object to store context for the current build
 // this is used to pass data between different steps of the build without having
 // to pass it through function arguments.
 // Not exhaustive, but should be extended to as needed whilst refactoring
 export const NextBuildContext: Partial<{
   compilerIdx?: number
-  serializedFlightMaps?: {
-    injectedClientEntries?: any
-    serverModuleIds?: any
-    edgeServerModuleIds?: any
-    asyncClientModules?: any
-    serverActions?: any
-    serverCSSManifest?: any
-    edgeServerCSSManifest?: any
-  }
+  pluginState: Record<string, any>
   serializedPagesManifestEntries: {
     edgeServerPages?: PagesManifest
     nodeServerPages?: PagesManifest

packages/next/src/build/webpack-build.ts

@@ -14,7 +14,11 @@ import * as Log from './output/log'
 import getBaseWebpackConfig, { loadProjectInfo } from './webpack-config'
 import { NextError } from '../lib/is-error'
 import { TelemetryPlugin } from './webpack/plugins/telemetry-plugin'
-import { NextBuildContext } from './build-context'
+import {
+  NextBuildContext,
+  resumePluginState,
+  getPluginState,
+} from './build-context'
 import { createEntrypoints } from './entries'
 import loadConfig from '../server/config'
 import { trace } from '../trace'
@@ -24,8 +28,6 @@ import {
   TurbotraceContext,
 } from './webpack/plugins/next-trace-entrypoints-plugin'
 import { UnwrapPromise } from '../lib/coalesced-function'
-import * as flightPluginModule from './webpack/plugins/flight-client-entry-plugin'
-import * as flightManifestPluginModule from './webpack/plugins/flight-manifest-plugin'
 import * as pagesPluginModule from './webpack/plugins/pages-manifest-plugin'
 import { Worker } from 'next/dist/compiled/jest-worker'
 import origDebug from 'next/dist/compiled/debug'
@@ -59,8 +61,8 @@ async function webpackBuildImpl(
   compilerName?: keyof typeof COMPILER_INDEXES
 ): Promise<{
   duration: number
+  pluginState: any
   turbotraceContext?: TurbotraceContext
-  serializedFlightMaps?: typeof NextBuildContext['serializedFlightMaps']
   serializedPagesManifestEntries?: typeof NextBuildContext['serializedPagesManifestEntries']
 }> {
   let result: CompilerResult | null = {
@@ -191,7 +193,9 @@ async function webpackBuildImpl(
 
     // Only continue if there were no errors
     if (!serverResult?.errors.length && !edgeServerResult?.errors.length) {
-      flightPluginModule.injectedClientEntries.forEach((value, key) => {
+      const pluginState = getPluginState()
+      for (const key in pluginState.injectedClientEntries) {
+        const value = pluginState.injectedClientEntries[key]
         const clientEntry = clientConfig.entry as webpack.EntryObject
         if (key === APP_CLIENT_INTERNALS) {
           clientEntry[CLIENT_STATIC_FILES_RUNTIME_MAIN_APP] = {
@@ -210,7 +214,7 @@ async function webpackBuildImpl(
             layer: WEBPACK_LAYERS.appClient,
           }
         }
-      })
+      }
 
       if (!compilerName || compilerName === 'client') {
         clientResult = await runCompiler(clientConfig, {
@@ -306,23 +310,7 @@ async function webpackBuildImpl(
     return {
       duration: webpackBuildEnd[0],
       turbotraceContext: traceEntryPointsPlugin?.turbotraceContext,
-      serializedFlightMaps: {
-        injectedClientEntries: Array.from(
-          flightPluginModule.injectedClientEntries.entries()
-        ),
-        serverModuleIds: Array.from(
-          flightPluginModule.serverModuleIds.entries()
-        ),
-        edgeServerModuleIds: Array.from(
-          flightPluginModule.edgeServerModuleIds.entries()
-        ),
-        asyncClientModules: Array.from(
-          flightManifestPluginModule.ASYNC_CLIENT_MODULES
-        ),
-        serverActions: flightPluginModule.serverActions,
-        serverCSSManifest: flightPluginModule.serverCSSManifest,
-        edgeServerCSSManifest: flightPluginModule.edgeServerCSSManifest,
-      },
+      pluginState: getPluginState(),
       serializedPagesManifestEntries: {
         edgeServerPages: pagesPluginModule.edgeServerPages,
         edgeServerAppPaths: pagesPluginModule.edgeServerAppPaths,
@@ -341,9 +329,11 @@ export async function workerMain(workerData: {
   // setup new build context from the serialized data passed from the parent
   Object.assign(NextBuildContext, workerData.buildContext)
 
+  // Resume plugin state
+  resumePluginState(NextBuildContext.pluginState)
+
   // restore module scope maps for flight plugins
-  const { serializedFlightMaps, serializedPagesManifestEntries } =
-    NextBuildContext
+  const { serializedPagesManifestEntries } = NextBuildContext
 
   for (const key of Object.keys(serializedPagesManifestEntries || {})) {
     Object.assign(
@@ -352,33 +342,6 @@ export async function workerMain(workerData: {
     )
   }
 
-  if (serializedFlightMaps) {
-    serializedFlightMaps.asyncClientModules?.forEach((item: any) =>
-      flightManifestPluginModule.ASYNC_CLIENT_MODULES.add(item)
-    )
-
-    for (const field of [
-      'injectedClientEntries',
-      'serverModuleIds',
-      'edgeServerModuleIds',
-    ]) {
-      for (const [key, value] of (serializedFlightMaps as any)[field] || []) {
-        ;(flightPluginModule as any)[field].set(key, value)
-      }
-    }
-
-    for (const field of [
-      'serverActions',
-      'serverCSSManifest',
-      'edgeServerCSSManifest',
-    ]) {
-      Object.assign(
-        (flightPluginModule as any)[field],
-        (serializedFlightMaps as any)[field]
-      )
-    }
-  }
-
   /// load the config because it's not serializable
   NextBuildContext.config = await loadConfig(
     PHASE_PRODUCTION_BUILD,
@@ -465,37 +428,9 @@ async function webpackBuildWithWorker() {
     // destroy worker so it's not sticking around using memory
     await worker.end()
 
-    prunedBuildContext.serializedFlightMaps = {
-      injectedClientEntries: [
-        ...(prunedBuildContext.serializedFlightMaps?.injectedClientEntries ||
-          []),
-        ...curResult.serializedFlightMaps?.injectedClientEntries,
-      ],
-      serverModuleIds: [
-        ...(prunedBuildContext.serializedFlightMaps?.serverModuleIds || []),
-        ...curResult.serializedFlightMaps?.serverModuleIds,
-      ],
-      edgeServerModuleIds: [
-        ...(prunedBuildContext.serializedFlightMaps?.edgeServerModuleIds || []),
-        ...curResult.serializedFlightMaps?.edgeServerModuleIds,
-      ],
-      asyncClientModules: [
-        ...(prunedBuildContext.serializedFlightMaps?.asyncClientModules || []),
-        ...curResult.serializedFlightMaps?.asyncClientModules,
-      ],
-      serverActions: {
-        ...prunedBuildContext.serializedFlightMaps?.serverActions,
-        ...curResult.serializedFlightMaps?.serverActions,
-      },
-      serverCSSManifest: {
-        ...prunedBuildContext.serializedFlightMaps?.serverCSSManifest,
-        ...curResult.serializedFlightMaps?.serverCSSManifest,
-      },
-      edgeServerCSSManifest: {
-        ...prunedBuildContext.serializedFlightMaps?.edgeServerCSSManifest,
-        ...curResult.serializedFlightMaps?.edgeServerCSSManifest,
-      },
-    }
+    // Update plugin state
+    prunedBuildContext.pluginState = curResult.pluginState
+
     prunedBuildContext.serializedPagesManifestEntries = {
       edgeServerAppPaths:
         curResult.serializedPagesManifestEntries?.edgeServerAppPaths,

packages/next/src/build/webpack/plugins/flight-client-entry-plugin.ts

@@ -22,7 +22,6 @@ import {
   SERVER_REFERENCE_MANIFEST,
   FLIGHT_SERVER_CSS_MANIFEST,
 } from '../../../shared/lib/constants'
-import { ASYNC_CLIENT_MODULES } from './flight-manifest-plugin'
 import {
   generateActionId,
   getActions,
@@ -32,6 +31,7 @@ import {
 import { traverseModules } from '../utils'
 import { normalizePathSep } from '../../../shared/lib/page-path/normalize-path-sep'
 import { isAppRouteRoute } from '../../../lib/is-app-route-route'
+import { getProxiedPluginState } from '../../build-context'
 
 interface Options {
   dev: boolean
@@ -41,11 +41,6 @@ interface Options {
 
 const PLUGIN_NAME = 'ClientEntryPlugin'
 
-export const injectedClientEntries = new Map<string, string>()
-
-export const serverModuleIds = new Map<string, string | number>()
-export const edgeServerModuleIds = new Map<string, string | number>()
-
 export type ActionManifest = {
   [actionId: string]: {
     workers: {
@@ -54,10 +49,25 @@ export type ActionManifest = {
   }
 }
 
-// A map to track "action" -> "list of bundles".
-export const serverActions: ActionManifest = {}
-export const serverCSSManifest: FlightCSSManifest = {}
-export const edgeServerCSSManifest: FlightCSSManifest = {}
+const pluginState = getProxiedPluginState({
+  // A map to track "action" -> "list of bundles".
+  serverActions: {} as ActionManifest,
+
+  // Manifest of CSS entry files for server/edge server.
+  serverCSSManifest: {} as FlightCSSManifest,
+  edgeServerCSSManifest: {} as FlightCSSManifest,
+
+  // Mapping of resource path to module id for server/edge server.
+  serverModuleIds: {} as Record<string, string | number>,
+  edgeServerModuleIds: {} as Record<string, string | number>,
+
+  // Collect modules from server/edge compiler in client layer,
+  // and detect if it's been used, and mark it as `async: true` for react.
+  // So that react could unwrap the async module from promise and render module itself.
+  ASYNC_CLIENT_MODULES: [] as string[],
+
+  injectedClientEntries: {} as Record<string, string>,
+})
 
 export class FlightClientEntryPlugin {
   dev: boolean
@@ -110,12 +120,11 @@ export class FlightClientEntryPlugin {
           }
 
           if (this.isEdgeServer) {
-            edgeServerModuleIds.set(
-              ssrNamedModuleId.replace(/\/next\/dist\/esm\//, '/next/dist/'),
-              modId
-            )
+            pluginState.edgeServerModuleIds[
+              ssrNamedModuleId.replace(/\/next\/dist\/esm\//, '/next/dist/')
+            ] = modId
           } else {
-            serverModuleIds.set(ssrNamedModuleId, modId)
+            pluginState.serverModuleIds[ssrNamedModuleId] = modId
           }
         }
       }
@@ -125,7 +134,7 @@ export class FlightClientEntryPlugin {
         // Using the client layer module, which doesn't have `rsc` tag in buildInfo.
         if (mod.request && mod.resource && !mod.buildInfo.rsc) {
           if (compilation.moduleGraph.isAsync(mod)) {
-            ASYNC_CLIENT_MODULES.add(mod.resource)
+            pluginState.ASYNC_CLIENT_MODULES.push(mod.resource)
           }
         }
 
@@ -292,9 +301,9 @@ export class FlightClientEntryPlugin {
     compilation.hooks.afterOptimizeModules.tap(PLUGIN_NAME, () => {
       const cssImportsForChunk: Record<string, Set<string>> = {}
 
-      let cssManifest = this.isEdgeServer
-        ? edgeServerCSSManifest
-        : serverCSSManifest
+      const cssManifest = this.isEdgeServer
+        ? pluginState.edgeServerCSSManifest
+        : pluginState.serverCSSManifest
 
       function collectModule(entryName: string, mod: any) {
         const resource = mod.resource
@@ -401,11 +410,11 @@ export class FlightClientEntryPlugin {
       (assets: webpack.Compilation['assets']) => {
         const manifest = JSON.stringify(
           {
-            ...serverCSSManifest,
-            ...edgeServerCSSManifest,
+            ...pluginState.serverCSSManifest,
+            ...pluginState.edgeServerCSSManifest,
             __entry_css_mods__: {
-              ...serverCSSManifest.__entry_css_mods__,
-              ...edgeServerCSSManifest.__entry_css_mods__,
+              ...pluginState.serverCSSManifest.__entry_css_mods__,
+              ...pluginState.edgeServerCSSManifest.__entry_css_mods__,
             },
           },
           null,
@@ -636,7 +645,7 @@ export class FlightClientEntryPlugin {
         entryData.lastActiveTime = Date.now()
       }
     } else {
-      injectedClientEntries.set(bundlePath, clientLoader)
+      pluginState.injectedClientEntries[bundlePath] = clientLoader
     }
 
     // Inject the entry to the server compiler (__sc_client__).
@@ -689,12 +698,12 @@ export class FlightClientEntryPlugin {
     for (const [p, names] of actionsArray) {
       for (const name of names) {
         const id = generateActionId(p, name)
-        if (typeof serverActions[id] === 'undefined') {
-          serverActions[id] = {
+        if (typeof pluginState.serverActions[id] === 'undefined') {
+          pluginState.serverActions[id] = {
             workers: {},
           }
         }
-        serverActions[id].workers[bundlePath] = ''
+        pluginState.serverActions[id].workers[bundlePath] = ''
       }
     }
 
@@ -761,14 +770,18 @@ export class FlightClientEntryPlugin {
       }
     })
 
-    for (let id in serverActions) {
-      const action = serverActions[id]
+    for (let id in pluginState.serverActions) {
+      const action = pluginState.serverActions[id]
       for (let name in action.workers) {
         action.workers[name] = actionModId[name]
       }
     }
 
-    const json = JSON.stringify(serverActions, null, this.dev ? 2 : undefined)
+    const json = JSON.stringify(
+      pluginState.serverActions,
+      null,
+      this.dev ? 2 : undefined
+    )
     assets[SERVER_REFERENCE_MANIFEST + '.js'] = new sources.RawSource(
       'self.__RSC_SERVER_MANIFEST=' + json
     ) as unknown as webpack.sources.RawSource