perf(vm): don't prewarm modules the worker never requests (#11033)

Co-authored-by: Vladimir Sheremet <sleuths.slews0s@icloud.com>

Author: MarshallOfSound

packages/mocker/src/node/hoistMocks.ts

@@ -1,11 +1,14 @@
 import type {
+  ArrowFunctionExpression,
   AwaitExpression,
   CallExpression,
   ExportDefaultDeclaration,
   ExportNamedDeclaration,
   Expression,
+  FunctionExpression,
   Identifier,
   ImportDeclaration,
+  SpreadElement,
   VariableDeclaration,
 } from 'estree'
 import type { Rollup } from 'vite'
@@ -16,7 +19,16 @@ import MagicString from 'magic-string'
 import { relative } from 'pathe'
 import { esmWalker } from './esmWalker'
 
+export interface StaticMockCall {
+  method: string
+  specifier: string
+  hasFactory: boolean
+  /** the factory uses `importOriginal`/`importActual` */
+  factoryLoadsOriginal: boolean
+}
+
 export interface HoistMocksOptions {
+  onStaticMock?: (call: StaticMockCall) => void
   /**
    * List of modules that should always be imported before compiler hints.
    * @default 'vitest'
@@ -340,6 +352,22 @@ export function hoistMocks(
               `Cannot export the result of "${method}". Remove export declaration because "${method}" doesn\'t return anything.`,
             )
           }
+          if (options.onStaticMock) {
+            const specifier = getStaticSpecifier(node.arguments[0])
+            if (specifier != null) {
+              // anything but an inline function may still load the original
+              const factory = node.arguments[1]?.type === 'ArrowFunctionExpression' || node.arguments[1]?.type === 'FunctionExpression'
+                ? node.arguments[1] as Positioned<ArrowFunctionExpression | FunctionExpression>
+                : undefined
+              options.onStaticMock({
+                method: methodName,
+                specifier,
+                hasFactory: factory != null,
+                factoryLoadsOriginal: factory != null
+                  && (factory.params.length > 0 || code.slice(factory.start, factory.end).includes('importActual')),
+              })
+            }
+          }
           // rewrite vi.mock(import('..')) into vi.mock('..')
           if (
             node.type === 'CallExpression'
@@ -610,3 +638,18 @@ function createIndexLocationsMap(source: string): Map<number, { line: number; co
   }
   return map
 }
+
+function getStaticSpecifier(node: Expression | SpreadElement | undefined): string | undefined {
+  if (node?.type === 'AwaitExpression') {
+    node = node.argument
+  }
+  if (node?.type === 'ImportExpression') {
+    node = node.source
+  }
+  if (node?.type === 'Literal' && typeof node.value === 'string') {
+    return node.value
+  }
+  if (node?.type === 'TemplateLiteral' && node.expressions.length === 0) {
+    return node.quasis[0].value.cooked ?? undefined
+  }
+}

packages/mocker/src/node/hoistMocksPlugin.ts

@@ -1,6 +1,6 @@
 import type { SourceMap } from 'magic-string'
 import type { Plugin, Rollup } from 'vite'
-import type { HoistMocksOptions } from './hoistMocks'
+import type { HoistMocksOptions, StaticMockCall } from './hoistMocks'
 import { createFilter } from 'vite'
 import { cleanUrl } from '../utils'
 import { hoistMocks } from './hoistMocks'
@@ -46,6 +46,7 @@ export function hoistMocksPlugin(options: HoistMocksPluginOptions = {}): Plugin
       if (!filter(id)) {
         return
       }
+      const staticMocks: StaticMockCall[] = []
       const s = hoistMocks(code, id, this.parse, {
         regexpHoistable,
         hoistableMockMethodNames,
@@ -55,12 +56,19 @@ export function hoistMocksPlugin(options: HoistMocksPluginOptions = {}): Plugin
         root,
         getMap: () => this.getCombinedSourcemap(),
         ...options,
+        onStaticMock(call) {
+          staticMocks.push(call)
+          options.onStaticMock?.(call)
+        },
       })
-      if (s) {
-        return {
-          code: s.toString(),
-          map: s.generateMap({ hires: 'boundary', source: cleanUrl(id) }),
-        }
+      // vite keeps `meta` across re-transforms, so always reset it
+      if (!s) {
+        return { meta: { vitestStaticMocks: null } }
+      }
+      return {
+        code: s.toString(),
+        map: s.generateMap({ hires: 'boundary', source: cleanUrl(id) }),
+        meta: { vitestStaticMocks: staticMocks },
       }
     },
   }

packages/mocker/src/node/index.ts

@@ -3,6 +3,7 @@ export { automockModule } from './automock'
 export type { AutomockPluginOptions } from './automockPlugin'
 export { automockPlugin } from './automockPlugin'
 export { dynamicImportPlugin } from './dynamicImportPlugin'
+export type { HoistMocksOptions, StaticMockCall } from './hoistMocks'
 export { hoistMockAndResolve as hoistMocks, hoistMocksPlugin } from './hoistMocksPlugin'
 export type { HoistMocksPluginOptions, HoistMocksResult } from './hoistMocksPlugin'
 export { interceptorPlugin } from './interceptorPlugin'

packages/vitest/src/node/cache/fsModuleCache.ts

@@ -1,4 +1,5 @@
-import type { DevEnvironment } from 'vite'
+import type { StaticMockCall } from '@vitest/mocker/node'
+import type { DevEnvironment, TransformResult } from 'vite'
 import type { ModuleType, VitestFetchResult } from '../../types/general'
 import type { Vitest } from '../core'
 import type { ResolvedConfig } from '../types/config'
@@ -38,7 +39,7 @@ export class FileSystemModuleCache {
   private rootCache: string
   private metadataFilePath: string
 
-  private version = '1.0.0-beta.6'
+  private version = '1.0.0-beta.7'
   private fsCacheRoots = new WeakMap<ResolvedConfig, string>()
   private fsEnvironmentHashMap = new WeakMap<DevEnvironment, string>()
   private fsCacheKeyGenerators = new Set<CacheKeyIdGenerator>()
@@ -136,12 +137,16 @@ export class FileSystemModuleCache {
       importedUrls: meta.importedUrls,
       mappings: meta.mappings,
       moduleType: meta.moduleType,
+      deps: meta.deps,
+      dynamicDeps: meta.dynamicDeps,
+      staticMocks: meta.staticMocks,
     }
   }
 
   async saveCachedModule(
     cachedFilePath: string,
     fetchResult: VitestFetchResult,
+    transformResult: TransformResult | null,
     importedUrls: string[] = [],
     mappings: boolean = false,
   ): Promise<void> {
@@ -153,6 +158,9 @@ export class FileSystemModuleCache {
         importedUrls,
         mappings,
         moduleType: fetchResult.moduleType,
+        deps: transformResult?.deps,
+        dynamicDeps: transformResult?.dynamicDeps,
+        staticMocks: transformResult?.__vitestStaticMocks,
       } satisfies Omit<CachedInlineModuleMeta, 'code'>
       debugFs?.(`${c.yellow('[write]')} ${fetchResult.id} is cached in ${cachedFilePath}`)
       await atomicWriteFile(cachedFilePath, `${fetchResult.code}${cacheComment}${this.toBase64(result)}`)
@@ -406,6 +414,9 @@ export interface CachedInlineModuleMeta {
   mappings: boolean
   importedUrls: string[]
   moduleType?: ModuleType
+  deps?: string[]
+  dynamicDeps?: string[]
+  staticMocks?: StaticMockCall[] | null
 }
 
 /**

packages/vitest/src/node/environments/fetchModule.ts

@@ -1,4 +1,5 @@
 import type { Span } from '@opentelemetry/api'
+import type { StaticMockCall } from '@vitest/mocker/node'
 import type { DevEnvironment, EnvironmentModuleNode, Rollup, TransformResult } from 'vite'
 import type { FetchFunctionOptions, FetchResult } from 'vite/module-runner'
 import type { FetchCachedFileSystemResult, ModuleType, VitestFetchResult } from '../../types/general'
@@ -122,7 +123,7 @@ class ModuleFetcher {
       }
 
       const tmpFile = join(tmpDir, hash('sha1', result.id, 'hex'))
-      return this.cacheResult(result, tmpFile).then((result) => {
+      return this.cacheResult(result, tmpFile, transformResult).then((result) => {
         if (transformResult) {
           transformResult.__vitestTmp = tmpFile
         }
@@ -148,7 +149,13 @@ class ModuleFetcher {
     const map = moduleGraphModule.transformResult?.map
     const mappings = map && !('version' in map) && map.mappings === ''
 
-    const cachedResult = await this.cacheResult(result, cachePath, importedUrls, !!mappings)
+    const cachedResult = await this.cacheResult(
+      result,
+      cachePath,
+      moduleGraphModule.transformResult,
+      importedUrls,
+      !!mappings,
+    )
     // remember where the code is stored on disk so that repeat fetches and the
     // `fetchWarmModules` snapshot can point at it in this session already, not
     // only after the cache is read back in the next one
@@ -287,8 +294,11 @@ class ModuleFetcher {
       code: cachedModule.code,
       map,
       ssr: true,
+      deps: cachedModule.deps,
+      dynamicDeps: cachedModule.dynamicDeps,
       __vitestTmp: cachePath,
       __vitestModuleType: moduleType,
+      __vitestStaticMocks: cachedModule.staticMocks,
     }
 
     // we populate the module graph to make the watch mode work because it relies on importers
@@ -339,6 +349,10 @@ class ModuleFetcher {
     if ('code' in result) {
       result.moduleType = await this.cachedModuleType(result.file, result.code, moduleGraphModule.transformResult)
     }
+    const transformResult = moduleGraphModule.transformResult
+    if (transformResult && moduleGraphModule.id) {
+      transformResult.__vitestStaticMocks ??= environment.pluginContainer.getModuleInfo(moduleGraphModule.id)?.meta?.vitestStaticMocks ?? null
+    }
     return result
   }
 
@@ -374,6 +388,7 @@ class ModuleFetcher {
   private async cacheResult(
     result: FetchResult,
     cachePath: string,
+    transformResult: TransformResult | null,
     importedUrls: string[] = [],
     mappings = false,
   ): Promise<FetchResult | FetchCachedFileSystemResult> {
@@ -386,7 +401,7 @@ class ModuleFetcher {
     }
 
     const savePromise = this.fsCache
-      .saveCachedModule(cachePath, result, importedUrls, mappings)
+      .saveCachedModule(cachePath, result, transformResult, importedUrls, mappings)
       .then(() => returnResult)
       .catch((error) => {
         debugFs?.(`failed to cache ${cachePath}, serving it inline: ${error}`)
@@ -588,5 +603,7 @@ declare module 'vite' {
     // `experimental.fsModuleCache` store or the forks pool's tmp copies
     __vitestTmp?: string
     __vitestModuleType?: ModuleType
+    // set by the hoistMocks plugin; null when the file was not hoisted
+    __vitestStaticMocks?: StaticMockCall[] | null
   }
 }