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) {