Skip to content

Commit b144ab5

Browse files
sheremet-vasheremet-va
authored
perf: serve warm modules to workers in one round-trip, enable the node compile cache (#10708)
1 parent 826714e commit b144ab5

8 files changed

Lines changed: 280 additions & 20 deletions

File tree

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

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
import type { DevEnvironment, FetchResult } from 'vite'
1+
import type { DevEnvironment } from 'vite'
2+
import type { ModuleType, VitestFetchResult } from '../../types/general'
23
import type { Vitest } from '../core'
34
import type { ResolvedConfig } from '../types/config'
45
import fs, { existsSync, mkdirSync, readFileSync } from 'node:fs'
@@ -33,7 +34,7 @@ export class FileSystemModuleCache {
3334
private rootCache: string
3435
private metadataFilePath: string
3536

36-
private version = '1.0.0-beta.4'
37+
private version = '1.0.0-beta.5'
3738
private fsCacheRoots = new WeakMap<ResolvedConfig, string>()
3839
private fsEnvironmentHashMap = new WeakMap<DevEnvironment, string>()
3940
private fsCacheKeyGenerators = new Set<CacheKeyIdGenerator>()
@@ -110,12 +111,13 @@ export class FileSystemModuleCache {
110111
code,
111112
importedUrls: meta.importedUrls,
112113
mappings: meta.mappings,
114+
moduleType: meta.moduleType,
113115
}
114116
}
115117

116-
async saveCachedModule<T extends FetchResult>(
118+
async saveCachedModule(
117119
cachedFilePath: string,
118-
fetchResult: T,
120+
fetchResult: VitestFetchResult,
119121
importedUrls: string[] = [],
120122
mappings: boolean = false,
121123
): Promise<void> {
@@ -126,6 +128,7 @@ export class FileSystemModuleCache {
126128
url: fetchResult.url,
127129
importedUrls,
128130
mappings,
131+
moduleType: fetchResult.moduleType,
129132
} satisfies Omit<CachedInlineModuleMeta, 'code'>
130133
debugFs?.(`${c.yellow('[write]')} ${fetchResult.id} is cached in ${cachedFilePath}`)
131134
await atomicWriteFile(cachedFilePath, `${fetchResult.code}${cacheComment}${this.toBase64(result)}`)
@@ -211,6 +214,7 @@ export class FileSystemModuleCache {
211214
mode: config.mode,
212215
consumer: config.consumer,
213216
resolve: config.resolve,
217+
injectCjsGlobal: vitestConfig.injectCjsGlobals,
214218
// plugins can have different options, so this is not the best key,
215219
// but we cannot access the options because there is no standard API for it
216220
plugins: config.plugins
@@ -366,6 +370,7 @@ export interface CachedInlineModuleMeta {
366370
code: string
367371
mappings: boolean
368372
importedUrls: string[]
373+
moduleType?: ModuleType
369374
}
370375

371376
/**

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

Lines changed: 42 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import type { Span } from '@opentelemetry/api'
22
import type { DevEnvironment, EnvironmentModuleNode, Rollup, TransformResult } from 'vite'
33
import type { FetchFunctionOptions, FetchResult } from 'vite/module-runner'
4-
import type { FetchCachedFileSystemResult, VitestFetchResult } from '../../types/general'
4+
import type { FetchCachedFileSystemResult, ModuleType, VitestFetchResult } from '../../types/general'
55
import type { OTELCarrier, Traces } from '../../utils/traces'
66
import type { FileSystemModuleCache } from '../cache/fsModuleCache'
77
import type { VitestResolver } from '../resolver'
@@ -142,7 +142,14 @@ class ModuleFetcher {
142142
const map = moduleGraphModule.transformResult?.map
143143
const mappings = map && !('version' in map) && map.mappings === ''
144144

145-
return this.cacheResult(result, cachePath, importedUrls, !!mappings)
145+
const cachedResult = await this.cacheResult(result, cachePath, importedUrls, !!mappings)
146+
// remember where the code is stored on disk so that repeat fetches and the
147+
// `fetchWarmModules` snapshot can point at it in this session already, not
148+
// only after the cache is read back in the next one
149+
if ('code' in result && moduleGraphModule.transformResult) {
150+
moduleGraphModule.transformResult.__vitestTmp = cachePath
151+
}
152+
return cachedResult
146153
}
147154

148155
// we need this for UI to be able to show a module graph
@@ -239,13 +246,11 @@ class ModuleFetcher {
239246
tmp: moduleGraphModule.transformResult.__vitestTmp,
240247
url: moduleGraphModule.url,
241248
invalidate: false,
242-
moduleType: this.detectModuleType
243-
? await detectModuleType(
244-
moduleGraphModule.file,
245-
moduleGraphModule.transformResult.code,
246-
this.sourceLoader(moduleGraphModule.file),
247-
)
248-
: undefined,
249+
moduleType: await this.cachedModuleType(
250+
moduleGraphModule.file,
251+
moduleGraphModule.transformResult.code,
252+
moduleGraphModule.transformResult,
253+
),
249254
}
250255
}
251256

@@ -264,11 +269,13 @@ class ModuleFetcher {
264269
if (!map && cachedModule.mappings) {
265270
map = { mappings: '' }
266271
}
272+
const moduleType = cachedModule.moduleType
267273
moduleGraphModule.transformResult = {
268274
code: cachedModule.code,
269275
map,
270276
ssr: true,
271277
__vitestTmp: cachePath,
278+
__vitestModuleType: moduleType,
272279
}
273280

274281
// we populate the module graph to make the watch mode work because it relies on importers
@@ -294,9 +301,7 @@ class ModuleFetcher {
294301
tmp: cachePath,
295302
url: cachedModule.url,
296303
invalidate: false,
297-
moduleType: this.detectModuleType
298-
? await detectModuleType(cachedModule.file, cachedModule.code, this.sourceLoader(cachedModule.file))
299-
: undefined,
304+
moduleType,
300305
}
301306
}
302307

@@ -318,8 +323,8 @@ class ModuleFetcher {
318323
).catch(handleRollupError)
319324

320325
const result: VitestFetchResult = processResultSource(environment, moduleRunnerModule)
321-
if (this.detectModuleType && 'code' in result) {
322-
result.moduleType = await detectModuleType(result.file, result.code, this.sourceLoader(result.file))
326+
if ('code' in result) {
327+
result.moduleType = await this.cachedModuleType(result.file, result.code, moduleGraphModule.transformResult)
323328
}
324329
return result
325330
}
@@ -331,6 +336,28 @@ class ModuleFetcher {
331336
return () => this.readFileConcurrently(file)
332337
}
333338

339+
// the module type is a pure function of the module, so detect it at most once
340+
// and memoize the verdict on the transform result. repeat fetches, the on-disk
341+
// cache (`cached`), and the `fetchWarmModules` snapshot all reuse it instead of
342+
// re-detecting. a no-op unless `injectCjsGlobals` is disabled — otherwise every
343+
// module receives the CJS globals and the type is irrelevant.
344+
private async cachedModuleType(
345+
file: string | null,
346+
code: string,
347+
transformResult: TransformResult | null | undefined,
348+
): Promise<ModuleType | undefined> {
349+
if (!this.detectModuleType) {
350+
return undefined
351+
}
352+
const moduleType
353+
= transformResult?.__vitestModuleType
354+
?? await detectModuleType(file, code, this.sourceLoader(file))
355+
if (transformResult) {
356+
transformResult.__vitestModuleType = moduleType
357+
}
358+
return moduleType
359+
}
360+
334361
private async cacheResult(
335362
result: FetchResult,
336363
cachePath: string,
@@ -543,5 +570,6 @@ export function handleRollupError(e: unknown): never {
543570
declare module 'vite' {
544571
export interface TransformResult {
545572
__vitestTmp?: string
573+
__vitestModuleType?: ModuleType
546574
}
547575
}

‎packages/vitest/src/node/pool.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,17 @@ export function createPool(ctx: Vitest): ProcessPool {
127127
...project.config.env,
128128
}
129129

130+
// V8 serializes compile-cached scripts without the source positions
131+
// that precise coverage relies on, so the compile cache must stay off
132+
// for the v8 provider (and custom providers, whose mechanism we can't
133+
// assume) in workers and any process they spawn. istanbul instruments
134+
// the source at transform time, so the cache is harmless there and the
135+
// boot speedup is kept.
136+
if (ctx.config.coverage.enabled && ctx.config.coverage.provider !== 'istanbul') {
137+
delete env.NODE_COMPILE_CACHE
138+
env.NODE_DISABLE_COMPILE_CACHE = '1'
139+
}
140+
130141
// env are case-insensitive on Windows, but spawned processes don't support it
131142
if (isWindows) {
132143
for (const name in env) {

‎packages/vitest/src/node/pools/rpc.ts‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import type { DevEnvironment, EnvironmentModuleNode, FetchResult } from 'vite'
2+
import type { FetchCachedFileSystemResult } from '../../types/general'
13
import type { RuntimeRPC } from '../../types/rpc'
24
import type { TestProject } from '../project'
35
import type { ResolveSnapshotPathHandlerContext } from '../types/config'
@@ -14,6 +16,19 @@ interface MethodsOptions {
1416
collect?: boolean
1517
}
1618

19+
// externalize verdicts served during this session, shared with fresh workers
20+
// via `fetchWarmModules`. Only verdicts for already-resolved urls are stored:
21+
// an unresolved specifier (a runtime-variable dynamic import of a bare name)
22+
// resolves through the requesting environment's plugin container, so its
23+
// verdict is importer-specific and cannot be shared.
24+
// Keyed by the DevEnvironment, not the server: a leading-slash url still
25+
// resolves to its id through that environment's plugin container, so a plugin
26+
// that resolves conditionally (e.g. on `this.environment`) can externalize the
27+
// same url in one environment and inline it in another — sharing the verdict
28+
// across environments would serve the wrong one. Per-environment keying also
29+
// drops the verdicts on a server restart, since environments are recreated.
30+
const warmExternals = new WeakMap<DevEnvironment, Record<string, FetchResult>>()
31+
1732
export function createMethodsRPC(project: TestProject, methodsOptions: MethodsOptions = {}): RuntimeRPC {
1833
const vitest = project.vitest
1934
const cacheFs = methodsOptions.cacheFs ?? false
@@ -47,6 +62,16 @@ export function createMethodsRPC(project: TestProject, methodsOptions: MethodsOp
4762
const metadata = project.vitest.state.metadata[project.name]
4863
if ('externalize' in result) {
4964
metadata.externalized[url] = result.externalize
65+
// builtins and network urls are already resolved inside the worker
66+
// without a round-trip, only module externalizations are worth sharing
67+
if (result.type === 'module' && url[0] === '/') {
68+
let externals = warmExternals.get(environment)
69+
if (!externals) {
70+
externals = Object.create(null) as Record<string, FetchResult>
71+
warmExternals.set(environment, externals)
72+
}
73+
externals[url] = result
74+
}
5075
}
5176
if ('tmp' in result) {
5277
metadata.tmps[url] = result.tmp
@@ -56,6 +81,75 @@ export function createMethodsRPC(project: TestProject, methodsOptions: MethodsOp
5681
return result
5782
})
5883
},
84+
async fetchWarmModules(environmentName, files) {
85+
const environment = project.vite.environments[environmentName]
86+
if (!environment) {
87+
throw new Error(`The environment ${environmentName} was not defined in the Vite config.`)
88+
}
89+
90+
const warm: Record<string, FetchResult | FetchCachedFileSystemResult> = Object.create(null)
91+
92+
// walk the import graphs of the requested files instead of dumping the
93+
// whole module graph — in large (watch) sessions the graph accumulates
94+
// modules this worker will never load
95+
const moduleGraph = environment.moduleGraph
96+
const queue: EnvironmentModuleNode[] = []
97+
for (const file of [...files, ...project.config.setupFiles]) {
98+
const nodes = moduleGraph.getModulesByFile(file)
99+
if (nodes) {
100+
queue.push(...nodes)
101+
}
102+
}
103+
104+
const seen = new Set<EnvironmentModuleNode>()
105+
while (queue.length) {
106+
const node = queue.pop()!
107+
if (seen.has(node)) {
108+
continue
109+
}
110+
seen.add(node)
111+
queue.push(...node.importedModules)
112+
113+
const transformResult = node.transformResult
114+
if (!transformResult || node.id == null) {
115+
continue
116+
}
117+
// the transformed code is already stored on disk either by the forks
118+
// pool (`cacheFs`) or by `experimental.fsModuleCache` — the worker can
119+
// read the file itself instead of fetching each module separately.
120+
// invalidated modules lose `transformResult` and drop out automatically
121+
const tmp = transformResult.__vitestTmp ?? (transformResult as { _vitest_tmp?: string })._vitest_tmp
122+
if (typeof tmp !== 'string') {
123+
continue
124+
}
125+
const entry: FetchCachedFileSystemResult = {
126+
cached: true,
127+
file: node.file,
128+
id: node.id,
129+
tmp,
130+
url: node.url,
131+
invalidate: false,
132+
// the fetch that stored this module on disk also memoized its module
133+
// type on the transform result (only when `injectCjsGlobals` is
134+
// disabled); reuse it so the evaluator injects the CJS globals for the
135+
// same modules it would on the direct-fetch path, no re-detection here
136+
moduleType: transformResult.__vitestModuleType,
137+
}
138+
warm[node.url] = entry
139+
if (node.id !== node.url) {
140+
warm[node.id] = entry
141+
}
142+
}
143+
144+
const externals = warmExternals.get(environment)
145+
if (externals) {
146+
for (const url in externals) {
147+
warm[url] ??= externals[url]
148+
}
149+
}
150+
151+
return warm
152+
},
59153
async resolve(id, importer, environmentName) {
60154
const environment = project.vite.environments[environmentName]
61155
if (!environment) {

‎packages/vitest/src/runtime/moduleRunner/startVitestModuleRunner.ts‎

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import type vm from 'node:vm'
2-
import type { EvaluatedModules } from 'vite/module-runner'
2+
import type { EvaluatedModules, FetchResult } from 'vite/module-runner'
3+
import type { FetchCachedFileSystemResult } from '../../types/general'
34
import type { WorkerGlobalState } from '../../types/worker'
45
import type { Traces } from '../../utils/traces'
56
import type { ExternalModulesExecutor } from '../external-executor'
@@ -54,6 +55,29 @@ export function startVitestModuleRunner(options: ContextModuleRunnerOptions): Vi
5455
}
5556
: undefined
5657

58+
// A fresh worker pays one strictly sequential `fetch` round-trip per module
59+
// in its test files' import graphs, even when the server processed all of
60+
// them already. Ask the server ONCE per run request for everything it has on
61+
// disk and answer those fetches locally. A file change invalidates the module
62+
// server-side, dropping it from the snapshot of every subsequent run request,
63+
// which keeps reused (isolate: false) workers in sync; an edit DURING a run
64+
// was racy before this fast path existed and stays racy with it — the
65+
// scheduled rerun always sees the fresh transform.
66+
let warmModules: Promise<Record<string, FetchResult | FetchCachedFileSystemResult> | null> | undefined
67+
let warmModulesContext: unknown
68+
69+
function fetchWarmModules() {
70+
const workerState = state()
71+
if (warmModulesContext !== workerState.ctx) {
72+
warmModulesContext = workerState.ctx
73+
warmModules = rpc()
74+
.fetchWarmModules(environment(), workerState.ctx.files.map(file => file.filepath))
75+
// if the snapshot cannot be fetched, fall back to per-module fetches
76+
.catch(() => null)
77+
}
78+
return warmModules!
79+
}
80+
5781
const evaluator = options.evaluator || new VitestModuleEvaluator(
5882
vm,
5983
{
@@ -144,6 +168,37 @@ export function startVitestModuleRunner(options: ContextModuleRunnerOptions): Vi
144168
return { cache: true }
145169
}
146170

171+
// only dependency fetches consult the snapshot: by the time the
172+
// first dependency is requested, the entry file is transformed and
173+
// its import graph is connected on the server, so the snapshot
174+
// actually covers the file's transitive dependencies
175+
if (importer != null) {
176+
const warm = await fetchWarmModules()
177+
// the null prototype is not preserved by the IPC serialization, so
178+
// ids like "constructor" must not fall through to Object.prototype
179+
const warmResult = warm && (
180+
Object.hasOwn(warm, id)
181+
? warm[id]
182+
: Object.hasOwn(warm, rawId)
183+
? warm[rawId]
184+
: undefined
185+
)
186+
if (warmResult) {
187+
if ('tmp' in warmResult) {
188+
try {
189+
const code = readFileSync(warmResult.tmp, 'utf-8')
190+
return { code, ...warmResult }
191+
}
192+
catch {
193+
// the tmp file is gone — fall back to a live fetch
194+
}
195+
}
196+
else {
197+
return warmResult
198+
}
199+
}
200+
}
201+
147202
const otelCarrier = traces?.getContextCarrier()
148203
const result = await rpc().fetch(
149204
id,

‎packages/vitest/src/types/rpc.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,12 @@ export interface RuntimeRPC {
1313
otelCarrier?: OTELCarrier,
1414
) => Promise<FetchResult | FetchCachedFileSystemResult>
1515
resolve: (id: string, importer: string | undefined, environment: string) => Promise<ResolveFunctionResult | null>
16+
/**
17+
* Returns the modules of the given test files' import graphs that the server
18+
* has already processed, so a fresh worker can load them from disk without
19+
* paying a `fetch` round-trip per module.
20+
*/
21+
fetchWarmModules: (environment: string, files: string[]) => Promise<Record<string, FetchResult | FetchCachedFileSystemResult>>
1622
transform: (id: string) => Promise<{ code?: string }>
1723

1824
onUserConsoleLog: (log: UserConsoleLog) => void

0 commit comments

Comments
 (0)