Skip to content

Commit d8b040c

Browse files
sheremet-vasheremet-va
authored
perf: reuse compiled code across vm pool contexts and prewarm the module graph (#10744)
1 parent e0fadba commit d8b040c

15 files changed

Lines changed: 478 additions & 55 deletions

File tree

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,7 @@ export class Vitest {
255255
this._resolver,
256256
resolved,
257257
this._fsCache,
258+
this.state,
258259
this._traces,
259260
this._tmpDir,
260261
)

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

Lines changed: 35 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,18 @@ import { normalizeResolvedIdToUrl } from './normalizeUrl'
1818
const saveCachePromises = new Map<string, Promise<VitestFetchResult>>()
1919
const readFilePromises = new Map<string, Promise<string | null>>()
2020

21+
/**
22+
* Tracks the wall time during which at least one transform is running.
23+
* Durations of individual fetches cannot be summed instead: concurrent
24+
* fetches (parallel workers, the vm pool graph prewarm) all wait on the same
25+
* deduplicated in-flight transforms, so per-caller wall times overcount the
26+
* actual work by orders of magnitude.
27+
*/
28+
export interface TransformClock {
29+
transformStarted: () => void
30+
transformFinished: () => void
31+
}
32+
2133
class ModuleFetcher {
2234
private tmpDirectories = new Set<string>()
2335
private fsCacheEnabled: boolean
@@ -30,6 +42,7 @@ class ModuleFetcher {
3042
private resolver: VitestResolver,
3143
private config: ResolvedConfig,
3244
private fsCache: FileSystemModuleCache,
45+
private clock: TransformClock,
3346
private tmpProjectDir: string,
3447
) {
3548
this.fsCacheEnabled = config.fsModuleCache === true
@@ -312,21 +325,27 @@ class ModuleFetcher {
312325
moduleGraphModule: EnvironmentModuleNode,
313326
options?: FetchFunctionOptions,
314327
): Promise<VitestFetchResult> {
315-
const moduleRunnerModule = await fetchModule(
316-
environment,
317-
url,
318-
importer,
319-
{
320-
...options,
321-
inlineSourceMap: false,
322-
},
323-
).catch(handleRollupError)
324-
325-
const result: VitestFetchResult = processResultSource(environment, moduleRunnerModule)
326-
if ('code' in result) {
327-
result.moduleType = await this.cachedModuleType(result.file, result.code, moduleGraphModule.transformResult)
328+
this.clock.transformStarted()
329+
try {
330+
const moduleRunnerModule = await fetchModule(
331+
environment,
332+
url,
333+
importer,
334+
{
335+
...options,
336+
inlineSourceMap: false,
337+
},
338+
).catch(handleRollupError)
339+
340+
const result: VitestFetchResult = processResultSource(environment, moduleRunnerModule)
341+
if ('code' in result) {
342+
result.moduleType = await this.cachedModuleType(result.file, result.code, moduleGraphModule.transformResult)
343+
}
344+
return result
345+
}
346+
finally {
347+
this.clock.transformFinished()
328348
}
329-
return result
330349
}
331350

332351
private sourceLoader(file: string | null): (() => Promise<string | null>) | undefined {
@@ -415,10 +434,11 @@ export function createFetchModuleFunction(
415434
resolver: VitestResolver,
416435
config: ResolvedConfig,
417436
fsCache: FileSystemModuleCache,
437+
clock: TransformClock,
418438
traces: Traces,
419439
tmpProjectDir: string,
420440
): VitestFetchFunction {
421-
const fetcher = new ModuleFetcher(resolver, config, fsCache, tmpProjectDir)
441+
const fetcher = new ModuleFetcher(resolver, config, fsCache, clock, tmpProjectDir)
422442
return async (url, importer, environment, cacheFs, options, otelCarrier) => {
423443
await traces.waitInit()
424444
const context = otelCarrier

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

Lines changed: 118 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import type { DevEnvironment, EnvironmentModuleNode, FetchResult } from 'vite'
2+
import type { FetchFunctionOptions } from 'vite/module-runner'
23
import type { FetchCachedFileSystemResult } from '../../types/general'
34
import type { RuntimeRPC } from '../../types/rpc'
5+
import type { OTELCarrier } from '../../utils/traces'
46
import type { TestProject } from '../project'
57
import type { ResolveSnapshotPathHandlerContext } from '../types/config'
68
import { existsSync, mkdirSync } from 'node:fs'
@@ -41,6 +43,56 @@ export function createMethodsRPC(project: TestProject, methodsOptions: MethodsOp
4143
mkdirSync(project.config.dumpDir, { recursive: true })
4244
}
4345
project.vitest.state.metadata[project.name].dumpDir = project.config.dumpDir
46+
47+
function getEnvironment(environmentName: string): DevEnvironment {
48+
const environment = project.vite.environments[environmentName]
49+
if (!environment) {
50+
throw new Error(`The environment ${environmentName} was not defined in the Vite config.`)
51+
}
52+
return environment
53+
}
54+
55+
async function fetchModule(
56+
url: string,
57+
importer: string | undefined,
58+
environment: DevEnvironment,
59+
options?: FetchFunctionOptions,
60+
otelCarrier?: OTELCarrier,
61+
// per-module durations are only recorded for direct worker fetches: the
62+
// graph prewarm fetches whole levels concurrently, so its per-module wall
63+
// times measure the queue position, not the module's own transform cost
64+
accountModuleDuration = true,
65+
): Promise<FetchResult | FetchCachedFileSystemResult> {
66+
const state = project.vitest.state
67+
const start = performance.now()
68+
69+
return await project._fetcher(url, importer, environment, cacheFs, options, otelCarrier).then((result) => {
70+
const metadata = state.metadata[project.name]
71+
if ('externalize' in result) {
72+
metadata.externalized[url] = result.externalize
73+
// builtins and network urls are already resolved inside the worker
74+
// without a round-trip, only module externalizations are worth sharing
75+
if (result.type === 'module' && url[0] === '/') {
76+
let externals = warmExternals.get(environment)
77+
if (!externals) {
78+
externals = Object.create(null) as Record<string, FetchResult>
79+
warmExternals.set(environment, externals)
80+
}
81+
externals[url] = result
82+
}
83+
}
84+
if ('tmp' in result) {
85+
metadata.tmps[url] = result.tmp
86+
}
87+
if (accountModuleDuration) {
88+
const duration = performance.now() - start
89+
metadata.duration[url] ??= []
90+
metadata.duration[url].push(duration)
91+
}
92+
return result
93+
})
94+
}
95+
4496
return {
4597
async fetch(
4698
url,
@@ -49,37 +101,7 @@ export function createMethodsRPC(project: TestProject, methodsOptions: MethodsOp
49101
options,
50102
otelCarrier,
51103
) {
52-
const environment = project.vite.environments[environmentName]
53-
if (!environment) {
54-
throw new Error(`The environment ${environmentName} was not defined in the Vite config.`)
55-
}
56-
57-
const start = performance.now()
58-
59-
return await project._fetcher(url, importer, environment, cacheFs, options, otelCarrier).then((result) => {
60-
const duration = performance.now() - start
61-
project.vitest.state.transformTime += duration
62-
const metadata = project.vitest.state.metadata[project.name]
63-
if ('externalize' in result) {
64-
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-
}
75-
}
76-
if ('tmp' in result) {
77-
metadata.tmps[url] = result.tmp
78-
}
79-
metadata.duration[url] ??= []
80-
metadata.duration[url].push(duration)
81-
return result
82-
})
104+
return fetchModule(url, importer, getEnvironment(environmentName), options, otelCarrier)
83105
},
84106
async fetchWarmModules(environmentName, files) {
85107
const environment = project.vite.environments[environmentName]
@@ -150,6 +172,71 @@ export function createMethodsRPC(project: TestProject, methodsOptions: MethodsOp
150172

151173
return warm
152174
},
175+
async prewarmModuleGraph(environmentName, files) {
176+
const environment = getEnvironment(environmentName)
177+
const moduleGraph = environment.moduleGraph
178+
const seen = new Set<string>()
179+
180+
async function walkNode(node: EnvironmentModuleNode): Promise<void> {
181+
const children: Promise<void>[] = []
182+
for (const child of node.importedModules) {
183+
if (child.url == null || seen.has(child.url)) {
184+
continue
185+
}
186+
if (child.transformResult) {
187+
seen.add(child.url)
188+
children.push(walkNode(child))
189+
}
190+
else {
191+
children.push(fetchNode(child.url, node.id ?? undefined))
192+
}
193+
}
194+
if (children.length) {
195+
await Promise.all(children)
196+
}
197+
}
198+
199+
async function fetchNode(url: string, importer: string | undefined): Promise<void> {
200+
if (seen.has(url)) {
201+
return
202+
}
203+
seen.add(url)
204+
try {
205+
await fetchModule(url, importer, environment, undefined, undefined, false)
206+
}
207+
catch {
208+
// the worker's own fetch will surface the error with the proper
209+
// import context
210+
return
211+
}
212+
let node: EnvironmentModuleNode | undefined
213+
try {
214+
node = await moduleGraph.getModuleByUrl(url) ?? moduleGraph.getModuleById(url) ?? undefined
215+
}
216+
catch {
217+
node = moduleGraph.getModuleById(url) ?? undefined
218+
}
219+
if (node) {
220+
await walkNode(node)
221+
}
222+
}
223+
224+
await Promise.all([...files, ...project.config.setupFiles].map(async (file) => {
225+
const nodes = moduleGraph.getModulesByFile(file)
226+
if (nodes && nodes.size) {
227+
await Promise.all(Array.from(nodes, (node) => {
228+
if (node.transformResult) {
229+
seen.add(node.url)
230+
return walkNode(node)
231+
}
232+
return fetchNode(node.url, undefined)
233+
}))
234+
}
235+
else {
236+
await fetchNode(file, undefined)
237+
}
238+
}))
239+
},
153240
async resolve(id, importer, environmentName) {
154241
const environment = project.vite.environments[environmentName]
155242
if (!environment) {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ export class TestProject {
101101
this._resolver,
102102
this.config,
103103
this.vitest._fsCache,
104+
this.vitest.state,
104105
this.vitest._traces,
105106
this.tmpDir,
106107
)

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

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import type { File, FileSpecification, Task, TaskResultPack } from '../runtime/runner/types'
22
import type { AsyncLeak, UserConsoleLog } from '../types/general'
3+
import type { TransformClock } from './environments/fetchModule'
34
import type { TestProject } from './project'
45
import type { MergedBlobs } from './reporters/blob'
56
import type { OnUnhandledErrorCallback } from './types/config'
@@ -15,7 +16,7 @@ function isAggregateError(err: unknown): err is AggregateError {
1516
return err instanceof Error && 'errors' in err
1617
}
1718

18-
export class StateManager {
19+
export class StateManager implements TransformClock {
1920
filesMap: Map<string, File[]> = new Map()
2021
pathsSet: Set<string> = new Set()
2122
idMap: Map<string, Task> = new Map()
@@ -24,7 +25,29 @@ export class StateManager {
2425
leakSet: Set<AsyncLeak> = new Set()
2526
reportedTasksMap: WeakMap<Task, TestModule | TestCase | TestSuite> = new WeakMap()
2627
blobs?: MergedBlobs
28+
/**
29+
* Wall time during which the server's module transform pipeline was busy,
30+
* measured as the union of in-flight fetch intervals. Individual fetch
31+
* durations cannot be summed instead: concurrent fetches (parallel workers,
32+
* the vm pool graph prewarm) all wait on the same deduplicated in-flight
33+
* transforms, so per-caller wall times overcount the actual work by orders
34+
* of magnitude.
35+
*/
2736
transformTime = 0
37+
private _transformsInflight = 0
38+
private _transformsBusyStart = 0
39+
40+
transformStarted(): void {
41+
if (this._transformsInflight++ === 0) {
42+
this._transformsBusyStart = performance.now()
43+
}
44+
}
45+
46+
transformFinished(): void {
47+
if (--this._transformsInflight === 0) {
48+
this.transformTime += performance.now() - this._transformsBusyStart
49+
}
50+
}
2851

2952
metadata: Record<string, {
3053
externalized: Record<string, string>

0 commit comments

Comments
 (0)
Sponsor
SponsoredKunjungi sekarang
Promo