feat: abstract moduleGraph into ModuleExecutionEnvironment

vitejs · Mar 10, 2024 · 5f5e0ec · 5f5e0ec
1 parent 83068fe
commit 5f5e0ec
Show file tree

Hide file tree

Showing 21 changed files with 136 additions and 90 deletions.
diff --git a/packages/vite/src/node/optimizer/optimizer.ts b/packages/vite/src/node/optimizer/optimizer.ts
@@ -539,7 +539,7 @@ async function createDepsOptimizer(
       // Cached transform results have stale imports (resolved to
       // old locations) so they need to be invalidated before the page is
       // reloaded.
-      server.getModuleGraph('browser').invalidateAll()
+      server.environments.get('browser')?.moduleGraph.invalidateAll()
 
       server.hot.send({
         type: 'full-reload',

diff --git a/packages/vite/src/node/plugins/asset.ts b/packages/vite/src/node/plugins/asset.ts
@@ -199,7 +199,9 @@ export function assetPlugin(config: ResolvedConfig): Plugin {
       // Inherit HMR timestamp if this asset was invalidated
       if (server) {
         const environment = options?.environment ?? 'browser'
-        const mod = server.getModuleGraph(environment).getModuleById(id)
+        const mod = server.environments
+          .get(environment)
+          ?.moduleGraph.getModuleById(id)
         if (mod && mod.lastHMRTimestamp > 0) {
           url = injectQuery(url, `t=${mod.lastHMRTimestamp}`)
         }

diff --git a/packages/vite/src/node/plugins/css.ts b/packages/vite/src/node/plugins/css.ts
@@ -942,7 +942,7 @@ export function cssAnalysisPlugin(config: ResolvedConfig): Plugin {
       }
 
       const environment = options?.environment ?? 'browser'
-      const moduleGraph = server?.getModuleGraph(environment)
+      const moduleGraph = server?.environments.get(environment)?.moduleGraph
       const thisModule = moduleGraph?.getModuleById(id)
 
       // Handle CSS @import dependency HMR and other added modules via this.addWatchFile.

diff --git a/packages/vite/src/node/plugins/esbuild.ts b/packages/vite/src/node/plugins/esbuild.ts
@@ -484,7 +484,7 @@ async function reloadOnTsconfigChange(changedFile: string) {
 
     // clear module graph to remove code compiled with outdated config
     server.environments.forEach((environment) =>
-      server.getModuleGraph(environment).invalidateAll(),
+      environment.moduleGraph.invalidateAll(),
     )
 
     // reset tsconfck so that recompile works with up2date configs

diff --git a/packages/vite/src/node/plugins/importAnalysis.ts b/packages/vite/src/node/plugins/importAnalysis.ts
@@ -239,10 +239,10 @@ export function importAnalysisPlugin(config: ResolvedConfig): Plugin {
 
       const depsOptimizer = getDepsOptimizer(config, ssr)
 
-      const moduleGraph = server.getModuleGraph(environment)
+      const moduleGraph = server.environments.get(environment)!.moduleGraph
       // since we are already in the transform phase of the importer, it must
       // have been loaded so its entry is guaranteed in the module graph.
-      const importerModule = moduleGraph.getModuleById(importer)
+      const importerModule = moduleGraph?.getModuleById(importer)
       if (!importerModule) {
         // This request is no longer valid. It could happen for optimized deps
         // requests. A full reload is going to request this id again.
@@ -525,9 +525,9 @@ export function importAnalysisPlugin(config: ResolvedConfig): Plugin {
             // record as safe modules
             // safeModulesPath should not include the base prefix.
             // See https://github.com/vitejs/vite/issues/9438#issuecomment-1465270409
-            server
-              ?.getModuleGraph('browser')
-              .safeModulesPath.add(fsPathFromUrl(stripBase(url, base)))
+            server?.browserEnvironment.moduleGraph.safeModulesPath.add(
+              fsPathFromUrl(stripBase(url, base)),
+            )
 
             if (url !== specifier) {
               let rewriteDone = false

diff --git a/packages/vite/src/node/plugins/importMetaGlob.ts b/packages/vite/src/node/plugins/importMetaGlob.ts
@@ -60,11 +60,11 @@ export function getAffectedGlobModules(
           (!negated.length || negated.every((glob) => isMatch(file, glob))),
       )
     ) {
-      const mod = server.getModuleGraph('browser').getModuleById(id)
+      const mod = server.browserEnvironment.moduleGraph.getModuleById(id)
 
       if (mod) {
         if (mod.file) {
-          server.getModuleGraph('browser').onFileChange(mod.file)
+          server.browserEnvironment.moduleGraph.onFileChange(mod.file)
         }
         modules.push(mod)
       }

diff --git a/packages/vite/src/node/plugins/worker.ts b/packages/vite/src/node/plugins/worker.ts
@@ -253,8 +253,9 @@ export function webWorkerPlugin(config: ResolvedConfig): Plugin {
             // dynamic worker type we can't know how import the env
             // so we copy /@vite/env code of server transform result into file header
             const environment = options?.environment ?? 'browser'
-            const moduleGraph = server.getModuleGraph(environment)
-            const module = moduleGraph.getModuleById(ENV_ENTRY)
+            const moduleGraph =
+              server.environments.get(environment)?.moduleGraph
+            const module = moduleGraph?.getModuleById(ENV_ENTRY)
             injectEnv = module?.transformResult?.code || ''
           }
         }

diff --git a/packages/vite/src/node/server/environment.ts b/packages/vite/src/node/server/environment.ts
@@ -0,0 +1,22 @@
+import type { PartialResolvedId } from 'rollup'
+import { EnvironmentModuleGraph } from './moduleGraph'
+
+export class ModuleExecutionEnvironment {
+  id: string
+  type: string
+  moduleGraph: EnvironmentModuleGraph
+  constructor(
+    id: string,
+    options: {
+      type: string
+      resolveId: (url: string) => Promise<PartialResolvedId | null>
+    },
+  ) {
+    this.id = id
+    this.type = options.type
+    this.moduleGraph = new EnvironmentModuleGraph(
+      options.type,
+      options.resolveId,
+    )
+  }
+}
diff --git a/packages/vite/src/node/server/hmr.ts b/packages/vite/src/node/server/hmr.ts
@@ -228,11 +228,11 @@ export async function handleHMRUpdate(
   // For now, we only call updateModules for the browser. Later on it should
   // also be called for each runtime.
 
-  let mods = server.getModuleGraph('browser').getModulesByFile(file)
+  let mods = server.browserEnvironment.moduleGraph.getModulesByFile(file)
   if (!mods) {
     // For now, given that the HMR SSR expects it, try to get the modules from the
     // server graph if the browser graph doesn't have it
-    mods = server.getModuleGraph('server').getModulesByFile(file)
+    mods = server.serverEnvironment.moduleGraph.getModulesByFile(file)
   }
 
   // check if any plugin wants to perform custom HMR handling
@@ -274,8 +274,8 @@ export async function handleHMRUpdate(
         hotContext.modules = filteredModules
           .map((mod) =>
             mod.id
-              ? server.getModuleGraph('browser').getModuleById(mod.id) ??
-                server.getModuleGraph('server').getModuleById(mod.id)
+              ? server.browserEnvironment.moduleGraph.getModuleById(mod.id) ??
+                server.serverEnvironment.moduleGraph.getModuleById(mod.id)
               : undefined,
           )
           .filter(Boolean) as EnvironmentModuleNode[]
@@ -327,9 +327,9 @@ export function updateModules(
 
     // TODO: we don't need NodeModule to have the runtime if we pass it to updateModules
     // it still seems useful to know the runtime for a given module
-    server
-      .getModuleGraph(mod.environment)
-      .invalidateModule(mod, invalidatedModules, timestamp, true)
+    server.environments
+      .get(mod.environment)
+      ?.moduleGraph.invalidateModule(mod, invalidatedModules, timestamp, true)
 
     if (needFullReload) {
       continue
@@ -424,9 +424,7 @@ export async function handleFileAddUnlink(
   isUnlink: boolean,
 ): Promise<void> {
   server.environments.forEach((environment) => {
-    const modules = [
-      ...(server.getModuleGraph(environment).getModulesByFile(file) || []),
-    ]
+    const modules = [...(environment.moduleGraph.getModulesByFile(file) || [])]
 
     if (isUnlink) {
       for (const deletedMod of modules) {

diff --git a/packages/vite/src/node/server/index.ts b/packages/vite/src/node/server/index.ts
@@ -72,7 +72,7 @@ import {
 } from './middlewares/static'
 import { timeMiddleware } from './middlewares/time'
 import type { EnvironmentModuleNode, ModuleNode } from './moduleGraph'
-import { EnvironmentModuleGraph, ModuleGraph } from './moduleGraph'
+import { ModuleGraph } from './moduleGraph'
 import { notFoundMiddleware } from './middlewares/notFound'
 import { errorMiddleware, prepareError } from './middlewares/error'
 import type { HMRBroadcaster, HmrOptions } from './hmr'
@@ -89,6 +89,7 @@ import type { TransformOptions, TransformResult } from './transformRequest'
 import { transformRequest } from './transformRequest'
 import { searchForWorkspaceRoot } from './searchRoot'
 import { warmupFiles } from './warmup'
+import { ModuleExecutionEnvironment } from './environment'
 
 export interface ServerOptions extends CommonServerOptions {
   /**
@@ -250,13 +251,14 @@ export interface ViteDevServer {
    */
   pluginContainer: PluginContainer
   /**
-   * Get the module graph for the given environment
+   * Dev Environments. Module execution environments attached to the Vite server.
    */
-  getModuleGraph(environment: string): EnvironmentModuleGraph
+  environments: Map<string, ModuleExecutionEnvironment>
   /**
-   * Available module graph environments
+   * Default environments
    */
-  environments: string[]
+  browserEnvironment: ModuleExecutionEnvironment
+  serverEnvironment: ModuleExecutionEnvironment
   /**
    * Module graph that tracks the import relationships, url to file mapping
    * and hmr state.
@@ -460,27 +462,32 @@ export async function _createServer(
       ) as FSWatcher)
     : createNoopWatcher(resolvedWatchOptions)
 
-  const browserModuleGraph = new EnvironmentModuleGraph('browser', (url) =>
-    container.resolveId(url, undefined, { ssr: false, environment: 'browser' }),
-  )
-  const serverModuleGraph = new EnvironmentModuleGraph('server', (url) =>
-    container.resolveId(url, undefined, { ssr: true, environment: 'server' }),
-  )
+  const browserEnvironment = new ModuleExecutionEnvironment('browser', {
+    type: 'browser',
+    resolveId: (url) =>
+      container.resolveId(url, undefined, {
+        ssr: false,
+        environment: 'browser',
+      }),
+  })
+  const serverEnvironment = new ModuleExecutionEnvironment('server', {
+    type: 'server',
+    resolveId: (url) =>
+      container.resolveId(url, undefined, { ssr: true, environment: 'server' }),
+  })
   const moduleGraph = new ModuleGraph({
-    browser: browserModuleGraph,
-    server: serverModuleGraph,
+    browser: browserEnvironment.moduleGraph,
+    server: serverEnvironment.moduleGraph,
   })
-  const environments = ['browser', 'server']
+  const environments = new Map([
+    ['browser', browserEnvironment],
+    ['server', serverEnvironment],
+  ])
+
+  // TODO: Later on, we may pass the environments map to pluginContainer instead of getModuleGraph
   const getModuleGraph = (environment: string) => {
-    if (environment === 'browser') {
-      return browserModuleGraph
-    } else if (environment === 'server') {
-      return serverModuleGraph
-    } else {
-      throw new Error(`Invalid module graph runtime: ${environment}`)
-    }
+    return environments.get(environment)!.moduleGraph
   }
-
   const container = await createPluginContainer(config, getModuleGraph, watcher)
   const closeHttpServer = createServerCloseFn(httpServer)
 
@@ -496,8 +503,9 @@ export async function _createServer(
     pluginContainer: container,
     ws,
     hot,
-    getModuleGraph,
     environments,
+    browserEnvironment,
+    serverEnvironment,
     moduleGraph,
     resolvedUrls: null, // will be set on listen
     ssrTransform(
@@ -747,12 +755,12 @@ export async function _createServer(
         publicFiles[isUnlink ? 'delete' : 'add'](path)
         if (!isUnlink) {
           const moduleWithSamePath =
-            await browserModuleGraph.getModuleByUrl(path)
+            await browserEnvironment.moduleGraph.getModuleByUrl(path)
           const etag = moduleWithSamePath?.transformResult?.etag
           if (etag) {
             // The public file should win on the next request over a module with the
             // same path. Prevent the transform etag fast path from serving the module
-            browserModuleGraph.etagToModuleMap.delete(etag)
+            browserEnvironment.moduleGraph.etagToModuleMap.delete(etag)
           }
         }
       }
@@ -766,7 +774,7 @@ export async function _createServer(
     await container.watchChange(file, { event: 'update' })
     // invalidate module graph cache on file change
     environments.forEach((environment) =>
-      getModuleGraph(environment).onFileChange(file),
+      environment.moduleGraph.onFileChange(file),
     )
     await onHMRUpdate(file, false)
   })
@@ -785,7 +793,9 @@ export async function _createServer(
     message?: string
     environment: string
   }) {
-    const mod = getModuleGraph(m.environment).urlToModuleMap.get(m.path)
+    const mod = environments
+      .get(m.environment)
+      ?.moduleGraph.urlToModuleMap.get(m.path)
     if (mod && mod.isSelfAccepting && mod.lastHMRTimestamp > 0) {
       config.logger.info(
         colors.yellow(`hmr invalidate `) +
@@ -809,7 +819,7 @@ export async function _createServer(
       invalidateModule({ path, message, environment })
     } else {
       environments.forEach((environment) => {
-        invalidateModule({ path, message, environment })
+        invalidateModule({ path, message, environment: environment.id })
       })
     }
   })

diff --git a/packages/vite/src/node/server/middlewares/indexHtml.ts b/packages/vite/src/node/server/middlewares/indexHtml.ts
@@ -127,7 +127,7 @@ const processNodeUrl = (
   // prefix with base (dev only, base is never relative)
   const replacer = (url: string) => {
     if (server) {
-      const mod = server.getModuleGraph('browser').urlToModuleMap.get(url)
+      const mod = server.browserEnvironment.moduleGraph.urlToModuleMap.get(url)
       if (mod && mod.lastHMRTimestamp > 0) {
         url = injectQuery(url, `t=${mod.lastHMRTimestamp}`)
       }
@@ -238,7 +238,7 @@ const devHtmlHook: IndexHtmlTransformHook = async (
     const modulePath = `${proxyModuleUrl}?html-proxy&index=${inlineModuleIndex}.${ext}`
 
     // invalidate the module so the newly cached contents will be served
-    const browserModuleGraph = server?.getModuleGraph('browser')
+    const browserModuleGraph = server?.browserEnvironment.moduleGraph
     const module = browserModuleGraph?.getModuleById(modulePath)
     if (module) {
       browserModuleGraph!.invalidateModule(module)
@@ -347,9 +347,11 @@ const devHtmlHook: IndexHtmlTransformHook = async (
       const url = `${proxyModulePath}?html-proxy&direct&index=${index}.css`
 
       // ensure module in graph after successful load
-      const mod = await server!
-        .getModuleGraph('browser')
-        .ensureEntryFromUrl(url, false)
+      const mod =
+        await server!.browserEnvironment.moduleGraph.ensureEntryFromUrl(
+          url,
+          false,
+        )
       ensureWatchedFile(watcher, mod.file, config.root)
 
       const result = await server!.pluginContainer.transform(code, mod.id!)
@@ -374,9 +376,11 @@ const devHtmlHook: IndexHtmlTransformHook = async (
       // will transform with css plugin and cache result with css-post plugin
       const url = `${proxyModulePath}?html-proxy&inline-css&style-attr&index=${index}.css`
 
-      const mod = await server!
-        .getModuleGraph('browser')
-        .ensureEntryFromUrl(url, false)
+      const mod =
+        await server!.browserEnvironment.moduleGraph.ensureEntryFromUrl(
+          url,
+          false,
+        )
       ensureWatchedFile(watcher, mod.file, config.root)
 
       await server?.pluginContainer.transform(code, mod.id!)

diff --git a/packages/vite/src/node/server/middlewares/static.ts b/packages/vite/src/node/server/middlewares/static.ts
@@ -217,7 +217,8 @@ export function isFileServingAllowed(
 
   if (server._fsDenyGlob(file)) return false
 
-  if (server.getModuleGraph('browser').safeModulesPath.has(file)) return true // TODO: safeModulePaths shouldn't be part of a moduleGraph
+  if (server.browserEnvironment.moduleGraph.safeModulesPath.has(file))
+    return true // TODO: safeModulePaths shouldn't be part of a moduleGraph
 
   if (
     server.config.server.fs.allow.some(

diff --git a/packages/vite/src/node/server/middlewares/transform.ts b/packages/vite/src/node/server/middlewares/transform.ts
@@ -50,9 +50,8 @@ export function cachedTransformMiddleware(
     // check if we can return 304 early
     const ifNoneMatch = req.headers['if-none-match']
     if (ifNoneMatch) {
-      const moduleByEtag = server
-        .getModuleGraph('browser')
-        .getModuleByEtag(ifNoneMatch)
+      const moduleByEtag =
+        server.browserEnvironment.moduleGraph.getModuleByEtag(ifNoneMatch)
       if (moduleByEtag?.transformResult?.etag === ifNoneMatch) {
         // For CSS requests, if the same CSS file is imported in a module,
         // the browser sends the request for the direct CSS request with the etag
@@ -143,7 +142,9 @@ export function transformMiddleware(
         } else {
           const originalUrl = url.replace(/\.map($|\?)/, '$1')
           const map = (
-            await server.getModuleGraph('browser').getModuleByUrl(originalUrl)
+            await server.browserEnvironment.moduleGraph.getModuleByUrl(
+              originalUrl,
+            )
           )?.transformResult?.map
           if (map) {
             return send(req, res, JSON.stringify(map), 'json', {
@@ -186,7 +187,7 @@ export function transformMiddleware(
           const ifNoneMatch = req.headers['if-none-match']
           if (
             ifNoneMatch &&
-            (await server.getModuleGraph('browser').getModuleByUrl(url))
+            (await server.browserEnvironment.moduleGraph.getModuleByUrl(url))
               ?.transformResult?.etag === ifNoneMatch
           ) {
             debugCache?.(`[304] ${prettifyUrl(url, server.config.root)}`)