feat: track module execution totalTime and selfTime (#8027)

Author: abrenneke

docs/advanced/api/test-module.md

@@ -95,5 +95,18 @@ interface ModuleDiagnostic {
    * This value is only available if the test was executed with `logHeapUsage` flag.
    */
   readonly heap: number | undefined
+  /**
+   * The time spent importing every non-externalized dependency that Vitest has processed.
+   */
+  readonly importDurations: Record<string, ImportDuration>
+}
+
+/** The time spent importing & executing a non-externalized file. */
+interface ImportDuration {
+  /** The time spent importing & executing the file itself, not counting all non-externalized imports that the file does. */
+  selfTime: number
+
+  /** The time spent importing & executing the file and all its imports. */
+  totalTime: number
 }
 ```

packages/runner/src/collect.ts

@@ -54,6 +54,11 @@ export async function collectTests(
 
       await runner.importFile(filepath, 'collect')
 
+      const durations = runner.getImportDurations?.()
+      if (durations) {
+        file.importDurations = durations
+      }
+
       const defaultTasks = await getDefaultSuite().collect(file)
 
       const fileHooks = createSuiteHooks()
@@ -85,6 +90,11 @@ export async function collectTests(
         state: 'fail',
         errors: [error],
       }
+
+      const durations = runner.getImportDurations?.()
+      if (durations) {
+        file.importDurations = durations
+      }
     }
 
     calculateSuiteHash(file)

packages/runner/src/types.ts

@@ -23,6 +23,7 @@ export type {
   Fixtures,
   HookCleanupCallback,
   HookListener,
+  ImportDuration,
   InferFixturesTypes,
   OnTestFailedHandler,
   OnTestFinishedHandler,

packages/runner/src/types/runner.ts

@@ -1,6 +1,7 @@
 import type { DiffOptions } from '@vitest/utils/diff'
 import type {
   File,
+  ImportDuration,
   SequenceHooks,
   SequenceSetupFiles,
   Suite,
@@ -154,6 +155,10 @@ export interface VitestRunner {
    * Function that is called when the runner attempts to get the value when `test.extend` is used with `{ injected: true }`
    */
   injectValue?: (key: string) => unknown
+  /**
+   * Gets the time spent importing each individual non-externalized file that Vitest collected.
+   */
+  getImportDurations?: () => Record<string, ImportDuration>
   /**
    * Publicly available configuration.
    */

packages/runner/src/types/tasks.ts

@@ -155,6 +155,15 @@ export interface TaskResult {
   pending?: boolean
 }
 
+/** The time spent importing & executing a non-externalized file. */
+export interface ImportDuration {
+  /** The time spent importing & executing the file itself, not counting all non-externalized imports that the file does. */
+  selfTime: number
+
+  /** The time spent importing & executing the file and all its imports. */
+  totalTime: number
+}
+
 /**
  * The tuple representing a single task update.
  * Usually reported after the task finishes.
@@ -239,6 +248,9 @@ export interface File extends Suite {
    * @internal
    */
   local?: boolean
+
+  /** The time spent importing every non-externalized dependency that Vitest has processed. */
+  importDurations?: Record<string, ImportDuration>
 }
 
 export interface Test<ExtraContext = object> extends TaskPopulated {

packages/vite-node/src/client.ts

@@ -176,7 +176,29 @@ export class ModuleCacheMap extends Map<string, ModuleCache> {
   }
 }
 
-export type ModuleExecutionInfo = Map<string, { startOffset: number }>
+export type ModuleExecutionInfo = Map<string, ModuleExecutionInfoEntry>
+
+export interface ModuleExecutionInfoEntry {
+  startOffset: number
+
+  /** The duration that was spent executing the module. */
+  duration: number
+
+  /** The time that was spent executing the module itself and externalized imports. */
+  selfTime: number
+}
+
+/** Stack to track nested module execution for self-time calculation. */
+type ExecutionStack = Array<{
+  /** The file that is being executed. */
+  filename: string
+
+  /** The start time of this module's execution. */
+  startTime: number
+
+  /** Accumulated time spent importing all sub-imports. */
+  subImportTime: number
+}>
 
 export class ViteNodeRunner {
   root: string
@@ -189,6 +211,27 @@ export class ViteNodeRunner {
    */
   moduleCache: ModuleCacheMap
 
+  /**
+   * Tracks the stack of modules being executed for the purpose of calculating import self-time.
+   *
+   * Note that while in most cases, imports are a linear stack of modules,
+   * this is occasionally not the case, for example when you have parallel top-level dynamic imports like so:
+   *
+   * ```ts
+   * await Promise.all([
+   *  import('./module1'),
+   *  import('./module2'),
+   * ]);
+   * ```
+   *
+   * In this case, the self time will be reported incorrectly for one of the modules (could go negative).
+   * As top-level awaits with dynamic imports like this are uncommon, we don't handle this case specifically.
+   */
+  private executionStack: ExecutionStack = []
+
+  // `performance` can be mocked, so make sure we're using the original function
+  private performanceNow = performance.now.bind(performance)
+
   constructor(public options: ViteNodeRunnerOptions) {
     this.root = options.root ?? process.cwd()
     this.moduleCache = options.moduleCache ?? new ModuleCacheMap()
@@ -542,10 +585,51 @@ export class ViteNodeRunner {
       columnOffset: -codeDefinition.length,
     }
 
-    this.options.moduleExecutionInfo?.set(options.filename, { startOffset: codeDefinition.length })
+    const finishModuleExecutionInfo = this.startCalculateModuleExecutionInfo(options.filename, codeDefinition.length)
 
-    const fn = vm.runInThisContext(code, options)
-    await fn(...Object.values(context))
+    try {
+      const fn = vm.runInThisContext(code, options)
+      await fn(...Object.values(context))
+    }
+    finally {
+      this.options.moduleExecutionInfo?.set(options.filename, finishModuleExecutionInfo())
+    }
+  }
+
+  /**
+   * Starts calculating the module execution info such as the total duration and self time spent on executing the module.
+   * Returns a function to call once the module has finished executing.
+   */
+  protected startCalculateModuleExecutionInfo(filename: string, startOffset: number): () => ModuleExecutionInfoEntry {
+    const startTime = this.performanceNow()
+
+    this.executionStack.push({
+      filename,
+      startTime,
+      subImportTime: 0,
+    })
+
+    return () => {
+      const duration = this.performanceNow() - startTime
+
+      const currentExecution = this.executionStack.pop()
+
+      if (currentExecution == null) {
+        throw new Error('Execution stack is empty, this should never happen')
+      }
+
+      const selfTime = duration - currentExecution.subImportTime
+
+      if (this.executionStack.length > 0) {
+        this.executionStack.at(-1)!.subImportTime += duration
+      }
+
+      return {
+        startOffset,
+        duration,
+        selfTime,
+      }
+    }
   }
 
   prepareContext(context: Record<string, any>): Record<string, any> {

packages/vitest/src/node/reporters/reported-tasks.ts

@@ -1,4 +1,5 @@
 import type {
+  ImportDuration,
   Task as RunnerTask,
   Test as RunnerTestCase,
   File as RunnerTestFile,
@@ -467,13 +468,15 @@ export class TestModule extends SuiteImplementation {
     const environmentSetupDuration = this.task.environmentLoad || 0
     const duration = this.task.result?.duration || 0
     const heap = this.task.result?.heap
+    const importDurations = this.task.importDurations ?? {}
     return {
       environmentSetupDuration,
       prepareDuration,
       collectDuration,
       setupDuration,
       duration,
       heap,
+      importDurations,
     }
   }
 }
@@ -626,6 +629,10 @@ export interface ModuleDiagnostic {
    * This value is only available if the test was executed with `logHeapUsage` flag.
    */
   readonly heap: number | undefined
+  /**
+   * The time spent importing every non-externalized dependency that Vitest has processed.
+   */
+  readonly importDurations: Record<string, ImportDuration>
 }
 
 function storeTask(

packages/vitest/src/public/index.ts

@@ -270,6 +270,7 @@ export type {
   ExtendedContext,
   HookCleanupCallback,
   HookListener,
+  ImportDuration,
   OnTestFailedHandler,
   OnTestFinishedHandler,
   RunMode,

packages/vitest/src/runtime/execute.ts

@@ -343,15 +343,20 @@ export class VitestExecutor extends ViteNodeRunner {
       columnOffset: -codeDefinition.length,
     }
 
-    this.options.moduleExecutionInfo?.set(options.filename, { startOffset: codeDefinition.length })
+    const finishModuleExecutionInfo = this.startCalculateModuleExecutionInfo(options.filename, codeDefinition.length)
 
-    const fn = vm.runInContext(code, vmContext, {
-      ...options,
-      // if we encountered an import, it's not inlined
-      importModuleDynamically: this.externalModules
-        .importModuleDynamically as any,
-    } as any)
-    await fn(...Object.values(context))
+    try {
+      const fn = vm.runInContext(code, vmContext, {
+        ...options,
+        // if we encountered an import, it's not inlined
+        importModuleDynamically: this.externalModules
+          .importModuleDynamically as any,
+      } as any)
+      await fn(...Object.values(context))
+    }
+    finally {
+      this.options.moduleExecutionInfo?.set(options.filename, finishModuleExecutionInfo())
+    }
   }
 
   public async importExternalModule(path: string): Promise<any> {

packages/vitest/src/runtime/runners/test.ts

@@ -2,6 +2,7 @@ import type { ExpectStatic } from '@vitest/expect'
 import type {
   CancelReason,
   File,
+  ImportDuration,
   Suite,
   Task,
   TestContext,
@@ -12,6 +13,7 @@ import type { SerializedConfig } from '../config'
 import type { VitestExecutor } from '../execute'
 import { getState, GLOBAL_EXPECT, setState } from '@vitest/expect'
 import { getNames, getTestName, getTests } from '@vitest/runner/utils'
+import { normalize } from 'pathe'
 import { createExpect } from '../../integrations/chai/index'
 import { inject } from '../../integrations/inject'
 import { getSnapshotClient } from '../../integrations/snapshot/chai'
@@ -181,6 +183,17 @@ export class VitestTestRunner implements VitestRunner {
     })
     return context
   }
+
+  getImportDurations(): Record<string, ImportDuration> {
+    const entries = [...(this.workerState.moduleExecutionInfo?.entries() ?? [])]
+    return Object.fromEntries(entries.map(([filepath, { duration, selfTime }]) => [
+      normalize(filepath),
+      {
+        selfTime,
+        totalTime: duration,
+      },
+    ]))
+  }
 }
 
 function clearModuleMocks(config: SerializedConfig) {

test/cli/test/reported-tasks.test.ts

@@ -276,6 +276,15 @@ it('correctly builds the full name', () => {
   expect(suiteSecondLevel.fullName).toBe('a group > a nested group')
 })
 
+it('correctly reports import durations', () => {
+  const diagnostic = testModule.diagnostic()
+
+  const filePath = resolve(root, './1_first.test.ts')
+  const importDuration = diagnostic.importDurations[filePath]
+  expect(importDuration.selfTime).toBeGreaterThan(0)
+  expect(importDuration.totalTime).toBeGreaterThan(0)
+})
+
 function date(time: Date) {
   return `${time.getDate()}/${time.getMonth() + 1}/${time.getFullYear()}`
 }

test/reporters/fixtures/import-durations-25ms-throws.ts

@@ -0,0 +1,5 @@
+await new Promise(resolve => setTimeout(resolve, 25))
+
+throw new Error('test')
+
+export {}

test/reporters/fixtures/import-durations-25ms.ts

@@ -0,0 +1,3 @@
+await new Promise(resolve => setTimeout(resolve, 25))
+
+export {}

test/reporters/fixtures/import-durations-50ms.ts

@@ -0,0 +1,3 @@
+import './import-durations-25ms'
+
+await new Promise(resolve => setTimeout(resolve, 50))

test/reporters/fixtures/import-durations-throws.test.ts

@@ -0,0 +1,6 @@
+import { expect, test } from 'vitest'
+import './import-durations-25ms-throws'
+
+test('works', async () => {
+  expect(true).toBe(true)
+})

test/reporters/fixtures/import-durations.test.ts

@@ -0,0 +1,6 @@
+import { expect, test } from 'vitest'
+import './import-durations-50ms'
+
+test('works', async () => {
+  expect(true).toBe(true)
+})

test/reporters/tests/__snapshots__/html.test.ts.snap

@@ -10,6 +10,7 @@ exports[`html reporter > resolves to "failing" status for test file "json-fail"
       "file": [Circular],
       "filepath": "<rootDir>/test/reporters/fixtures/json-fail.test.ts",
       "id": 0,
+      "importDurations": {},
       "meta": {},
       "mode": "run",
       "name": "json-fail.test.ts",
@@ -121,6 +122,7 @@ exports[`html reporter > resolves to "passing" status for test file "all-passing
       "file": [Circular],
       "filepath": "<rootDir>/test/reporters/fixtures/all-passing-or-skipped.test.ts",
       "id": 0,
+      "importDurations": {},
       "meta": {},
       "mode": "run",
       "name": "all-passing-or-skipped.test.ts",

test/reporters/tests/html.test.ts

@@ -27,6 +27,7 @@ describe('html reporter', async () => {
     file.environmentLoad = 0
     file.prepareDuration = 0
     file.setupDuration = 0
+    file.importDurations = {}
     file.result.duration = 0
     file.result.startTime = 0
     const task = file.tasks[0]
@@ -55,6 +56,7 @@ describe('html reporter', async () => {
     file.environmentLoad = 0
     file.prepareDuration = 0
     file.setupDuration = 0
+    file.importDurations = {}
     file.result.duration = 0
     file.result.startTime = 0
     const task = file.tasks[0]