feat: add caching to run failed and longer tests first (#1541)

* feat: add caching to run failed and longer tests first

* chore: cleanup

* chore: add clearCache

* chore: update lockfile

* chore: add filesCache

* refactor: add sequelizer

* chore: lockfile

* refactor: renaming

* chore: dont override version

* chore: fix clearCache

* chore: guard slice

* chore: cleanup

* docs: cleanup

* chore: cleanup

* chore: remove command for now

* refactor: cleanup

docs/config/index.md

@@ -565,3 +565,16 @@ RegExp pattern for files that will return en empty CSS file.
 A number of tests that are allowed to run at the same time marked with `test.concurrent`.
 
 Test above this limit will be queued to run when available slot appears.
+
+### cache
+
+- **Type**: `false | { dir? }`
+
+Options to configure Vitest cache policy. At the moment Vitest stores cache for test results to run the longer and failed tests first.
+
+#### cache.dir
+
+- **Type**: `string`
+- **Default**: `node_modules/.vitest`
+
+Path to cache directory.

docs/guide/cli.md

@@ -36,6 +36,10 @@ Useful to run with [`lint-staged`](https://github.com/okonet/lint-staged) or wit
 vitest related /src/index.ts /src/hello-world.js
 ```
 
+### `vitest clean cache`
+
+Clears cache folder.
+
 ## Options
 
 | Options       |               |

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

@@ -0,0 +1,23 @@
+import fs, { type Stats } from 'fs'
+
+type FileStatsCache = Pick<Stats, 'size'>
+
+export class FilesStatsCache {
+  public cache = new Map<string, FileStatsCache>()
+
+  public getStats(fsPath: string): FileStatsCache | undefined {
+    return this.cache.get(fsPath)
+  }
+
+  public async updateStats(fsPath: string) {
+    if (!fs.existsSync(fsPath))
+      return
+
+    const stats = await fs.promises.stat(fsPath)
+    this.cache.set(fsPath, { size: stats.size })
+  }
+
+  public removeStats(fsPath: string) {
+    this.cache.delete(fsPath)
+  }
+}

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

@@ -0,0 +1,38 @@
+import fs from 'fs'
+import { findUp } from 'find-up'
+import { resolve } from 'pathe'
+import { loadConfigFromFile } from 'vite'
+import { configFiles } from '../../constants'
+import type { CliOptions } from '../cli-api'
+import { slash } from '../../utils'
+
+export class VitestCache {
+  static resolveCacheDir(root: string, dir: string | undefined) {
+    return resolve(root, slash(dir || 'node_modules/.vitest'))
+  }
+
+  static async clearCache(options: CliOptions) {
+    const root = resolve(options.root || process.cwd())
+
+    const configPath = options.config
+      ? resolve(root, options.config)
+      : await findUp(configFiles, { cwd: root } as any)
+
+    const config = await loadConfigFromFile({ command: 'serve', mode: 'test' }, configPath)
+
+    const cache = config?.config.test?.cache
+
+    if (cache === false)
+      throw new Error('Cache is disabled')
+
+    const cachePath = VitestCache.resolveCacheDir(root, cache?.dir)
+
+    let cleared = false
+
+    if (fs.existsSync(cachePath)) {
+      fs.rmSync(cachePath, { recursive: true, force: true })
+      cleared = true
+    }
+    return { dir: cachePath, cleared }
+  }
+}

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

@@ -0,0 +1,76 @@
+import fs from 'fs'
+import { dirname, resolve } from 'pathe'
+import type { File, ResolvedConfig } from '../../types'
+import { version } from '../../../package.json'
+
+export interface SuiteResultCache {
+  failed: boolean
+  duration: number
+}
+
+export class ResultsCache {
+  private cache = new Map<string, SuiteResultCache>()
+  private cachePath: string | null = null
+  private version: string = version
+  private root = '/'
+
+  setConfig(root: string, config: ResolvedConfig['cache']) {
+    this.root = root
+    if (config)
+      this.cachePath = resolve(config.dir, 'results.json')
+  }
+
+  getResults(fsPath: string) {
+    return this.cache.get(fsPath?.slice(this.root.length))
+  }
+
+  async readFromCache() {
+    if (!this.cachePath)
+      return
+
+    if (fs.existsSync(this.cachePath)) {
+      const resultsCache = await fs.promises.readFile(this.cachePath, 'utf8')
+      const { results, version } = JSON.parse(resultsCache)
+      this.cache = new Map(results)
+      this.version = version
+    }
+  }
+
+  updateResults(files: File[]) {
+    files.forEach((file) => {
+      const result = file.result
+      if (!result)
+        return
+      const duration = result.duration || 0
+      // store as relative, so cache would be the same in CI and locally
+      const relativePath = file.filepath?.slice(this.root.length)
+      this.cache.set(relativePath, {
+        duration: duration >= 0 ? duration : 0,
+        failed: result.state === 'fail',
+      })
+    })
+  }
+
+  removeFromCache(filepath: string) {
+    this.cache.delete(filepath)
+  }
+
+  async writeToCache() {
+    if (!this.cachePath)
+      return
+
+    const results = Array.from(this.cache.entries())
+
+    const cacheDirname = dirname(this.cachePath)
+
+    if (!fs.existsSync(cacheDirname))
+      await fs.promises.mkdir(cacheDirname, { recursive: true })
+
+    const cache = JSON.stringify({
+      version: this.version,
+      results,
+    })
+
+    await fs.promises.writeFile(this.cachePath, cache)
+  }
+}

packages/vitest/src/node/config.ts

@@ -8,6 +8,7 @@ import { defaultPort } from '../constants'
 import { configDefaults } from '../defaults'
 import { resolveC8Options } from '../integrations/coverage'
 import { toArray } from '../utils'
+import { VitestCache } from './cache'
 
 const extraInlineDeps = [
   /^(?!.*(?:node_modules)).*\.mjs$/,
@@ -181,5 +182,9 @@ export function resolveConfig(
   if (typeof resolved.css === 'object')
     resolved.css.include ??= [/\.module\./]
 
+  resolved.cache ??= { dir: '' }
+  if (resolved.cache)
+    resolved.cache.dir = VitestCache.resolveCacheDir(resolved.root, resolved.cache.dir)
+
   return resolved
 }

packages/vitest/src/node/core.ts

@@ -91,6 +91,9 @@ export class Vitest {
 
     if (resolved.coverage.enabled)
       await cleanCoverage(resolved.coverage, resolved.coverage.clean)
+
+    this.state.results.setConfig(resolved.root, resolved.cache)
+    await this.state.results.readFromCache()
   }
 
   getSerializableConfig() {
@@ -133,6 +136,9 @@ export class Vitest {
       process.exit(exitCode)
     }
 
+    // populate once, update cache on watch
+    await Promise.all(files.map(file => this.state.stats.updateStats(file)))
+
     await this.runFiles(files)
 
     if (this.config.coverage.enabled)
@@ -205,7 +211,7 @@ export class Vitest {
     return runningTests
   }
 
-  async runFiles(files: string[]) {
+  async runFiles(paths: string[]) {
     await this.runningPromise
 
     this.runningPromise = (async () => {
@@ -217,16 +223,21 @@ export class Vitest {
       this.snapshot.clear()
       this.state.clearErrors()
       try {
-        await this.pool.runTests(files, invalidates)
+        await this.pool.runTests(paths, invalidates)
       }
       catch (err) {
         this.state.catchError(err, 'Unhandled Error')
       }
 
-      if (hasFailed(this.state.getFiles()))
+      const files = this.state.getFiles()
+
+      if (hasFailed(files))
         process.exitCode = 1
 
-      await this.report('onFinished', this.state.getFiles(), this.state.getUnhandledErrors())
+      await this.report('onFinished', files, this.state.getUnhandledErrors())
+
+      this.state.results.updateResults(files)
+      await this.state.results.writeToCache()
     })()
       .finally(() => {
         this.runningPromise = undefined
@@ -352,6 +363,8 @@ export class Vitest {
 
       if (this.state.filesMap.has(id)) {
         this.state.filesMap.delete(id)
+        this.state.results.removeFromCache(id)
+        this.state.stats.removeStats(id)
         this.changedTests.delete(id)
         this.report('onTestRemoved', id)
       }
@@ -360,6 +373,7 @@ export class Vitest {
       id = slash(id)
       if (await this.isTargetFile(id)) {
         this.changedTests.add(id)
+        await this.state.stats.updateStats(id)
         this.scheduleRerun(id)
       }
     }

packages/vitest/src/node/pool.ts

@@ -1,7 +1,6 @@
 import { MessageChannel } from 'worker_threads'
 import { pathToFileURL } from 'url'
 import { cpus } from 'os'
-import { createHash } from 'crypto'
 import { resolve } from 'pathe'
 import type { Options as TinypoolOptions } from 'tinypool'
 import { Tinypool } from 'tinypool'
@@ -9,8 +8,9 @@ import { createBirpc } from 'birpc'
 import type { RawSourceMap } from 'vite-node'
 import type { ResolvedConfig, WorkerContext, WorkerRPC } from '../types'
 import { distDir } from '../constants'
-import { AggregateError, slash } from '../utils'
+import { AggregateError } from '../utils'
 import type { Vitest } from './core'
+import { BaseSequelizer } from './sequelizers/BaseSequelizer'
 
 export type RunWithFiles = (files: string[], invalidates?: string[]) => Promise<void>
 
@@ -86,29 +86,15 @@ export function createPool(ctx: Vitest): WorkerPool {
       }
     }
 
+    const sequelizer = new BaseSequelizer(ctx)
+
     return async (files, invalidates) => {
       const config = ctx.getSerializableConfig()
 
-      if (config.shard) {
-        const { index, count } = config.shard
-        const shardSize = Math.ceil(files.length / count)
-        const shardStart = shardSize * (index - 1)
-        const shardEnd = shardSize * index
-        files = files
-          .map((file) => {
-            const fullPath = resolve(slash(config.root), slash(file))
-            const specPath = fullPath.slice(config.root.length)
-            return {
-              file,
-              hash: createHash('sha1')
-                .update(specPath)
-                .digest('hex'),
-            }
-          })
-          .sort((a, b) => (a.hash < b.hash ? -1 : a.hash > b.hash ? 1 : 0))
-          .slice(shardStart, shardEnd)
-          .map(({ file }) => file)
-      }
+      if (config.shard)
+        files = await sequelizer.shard(files)
+
+      files = await sequelizer.sort(files)
 
       if (!ctx.config.threads) {
         await runFiles(config, files)

packages/vitest/src/node/sequelizers/BaseSequelizer.ts

@@ -0,0 +1,66 @@
+import { createHash } from 'crypto'
+import { resolve } from 'pathe'
+import { slash } from 'vite-node/utils'
+import type { Vitest } from '../core'
+import type { TestSequelizer } from './types'
+
+export class BaseSequelizer implements TestSequelizer {
+  protected ctx: Vitest
+
+  constructor(ctx: Vitest) {
+    this.ctx = ctx
+  }
+
+  // async so it can be extended by other sequelizers
+  public async shard(files: string[]): Promise<string[]> {
+    const { config } = this.ctx
+    const { index, count } = config.shard!
+    const shardSize = Math.ceil(files.length / count)
+    const shardStart = shardSize * (index - 1)
+    const shardEnd = shardSize * index
+    return [...files]
+      .map((file) => {
+        const fullPath = resolve(slash(config.root), slash(file))
+        const specPath = fullPath?.slice(config.root.length)
+        return {
+          file,
+          hash: createHash('sha1')
+            .update(specPath)
+            .digest('hex'),
+        }
+      })
+      .sort((a, b) => (a.hash < b.hash ? -1 : a.hash > b.hash ? 1 : 0))
+      .slice(shardStart, shardEnd)
+      .map(({ file }) => file)
+  }
+
+  // async so it can be extended by other sequelizers
+  public async sort(files: string[]): Promise<string[]> {
+    const { state } = this.ctx
+    return [...files].sort((a, b) => {
+      const aState = state.getFileTestResults(a)
+      const bState = state.getFileTestResults(b)
+
+      if (!aState || !bState) {
+        const statsA = state.getFileStats(a)
+        const statsB = state.getFileStats(b)
+
+        // run unknown first
+        if (!statsA || !statsB)
+          return !statsA && statsB ? -1 : !statsB && statsA ? 1 : 0
+
+        // run larger files first
+        return statsB.size - statsA.size
+      }
+
+      // run failed first
+      if (aState.failed && !bState.failed)
+        return -1
+      if (!aState.failed && bState.failed)
+        return 1
+
+      // run longer first
+      return bState.duration - aState.duration
+    })
+  }
+}

packages/vitest/src/node/sequelizers/types.ts

@@ -0,0 +1,15 @@
+import type { Awaitable } from '../../types'
+import type { Vitest } from '../core'
+
+export interface TestSequelizer {
+  /**
+   * Slicing tests into shards. Will be run before `sort`.
+   * Only run, if `shard` is defined.
+   */
+  shard(files: string[]): Awaitable<string[]>
+  sort(files: string[]): Awaitable<string[]>
+}
+
+export interface TestSequelizerContructor {
+  new (ctx: Vitest): TestSequelizer
+}

packages/vitest/src/node/state.ts

@@ -1,10 +1,22 @@
 import type { ErrorWithDiff, File, Task, TaskResultPack, UserConsoleLog } from '../types'
+import { FilesStatsCache } from './cache/files'
+import { ResultsCache } from './cache/results'
 
 export class StateManager {
   filesMap = new Map<string, File>()
   idMap = new Map<string, Task>()
   taskFileMap = new WeakMap<Task, File>()
   errorsSet = new Set<unknown>()
+  results = new ResultsCache()
+  stats = new FilesStatsCache()
+
+  getFileTestResults(id: string) {
+    return this.results.getResults(id)
+  }
+
+  getFileStats(id: string) {
+    return this.stats.getStats(id)
+  }
 
   catchError(err: unknown, type: string) {
     (err as ErrorWithDiff).type = type

packages/vitest/src/types/config.ts

@@ -361,6 +361,14 @@ export interface InlineConfig {
    * @default 5
    */
   maxConcurrency?: number
+
+  /**
+   * Options for configuring cache policy.
+   * @default { dir: 'node_modules/.vitest' }
+   */
+  cache?: false | {
+    dir?: string
+  }
 }
 
 export interface UserConfig extends InlineConfig {
@@ -407,7 +415,7 @@ export interface UserConfig extends InlineConfig {
   shard?: string
 }
 
-export interface ResolvedConfig extends Omit<Required<UserConfig>, 'config' | 'filters' | 'coverage' | 'testNamePattern' | 'related' | 'api' | 'reporters' | 'resolveSnapshotPath' | 'shard'> {
+export interface ResolvedConfig extends Omit<Required<UserConfig>, 'config' | 'filters' | 'coverage' | 'testNamePattern' | 'related' | 'api' | 'reporters' | 'resolveSnapshotPath' | 'shard' | 'cache'> {
   base?: string
 
   config?: string
@@ -427,4 +435,8 @@ export interface ResolvedConfig extends Omit<Required<UserConfig>, 'config' | 'f
     index: number
     count: number
   }
+
+  cache: {
+    dir: string
+  } | false
 }

pnpm-lock.yaml

@@ -199,7 +199,7 @@ importers:
       typescript: 4.6.3
       vitest: workspace:*
     dependencies:
-      next: 12.1.5_ezdxe4hg7n3pawg24sxf3xmgta
+      next: 12.1.5_zpnidt7m3osuk7shl3s4oenomq
       react: 18.0.0
       react-dom: 18.0.0_react@18.0.0
     devDependencies:
@@ -768,6 +768,12 @@ importers:
     devDependencies:
       rollup: 2.75.7
 
+  test/cache:
+    specifiers:
+      vitest: workspace:*
+    devDependencies:
+      vitest: link:../../packages/vitest
+
   test/cjs:
     specifiers:
       '@types/fs-extra': ^9.0.13
@@ -15837,7 +15843,7 @@ packages:
     resolution: {integrity: sha512-9iN1ka/9zmX1ZvLV9ewJYEk9h7RyRRtqdK0woXcqohu8EWIerfPUjYJPg0ULy0UqP7cslmdGc8xKDJcojlKiaw==}
     dev: true
 
-  /next/12.1.5_ezdxe4hg7n3pawg24sxf3xmgta:
+  /next/12.1.5_zpnidt7m3osuk7shl3s4oenomq:
     resolution: {integrity: sha512-YGHDpyfgCfnT5GZObsKepmRnne7Kzp7nGrac07dikhutWQug7hHg85/+sPJ4ZW5Q2pDkb+n0FnmLkmd44htIJQ==}
     engines: {node: '>=12.22.0'}
     hasBin: true
@@ -15860,7 +15866,7 @@ packages:
       postcss: 8.4.5
       react: 18.0.0
       react-dom: 18.0.0_react@18.0.0
-      styled-jsx: 5.0.1_uyynoipo3v3vrfv6si7tyrw7ku
+      styled-jsx: 5.0.1_react@18.0.0
     optionalDependencies:
       '@next/swc-android-arm-eabi': 12.1.5
       '@next/swc-android-arm64': 12.1.5
@@ -19034,7 +19040,7 @@ packages:
       inline-style-parser: 0.1.1
     dev: true
 
-  /styled-jsx/5.0.1_uyynoipo3v3vrfv6si7tyrw7ku:
+  /styled-jsx/5.0.1_react@18.0.0:
     resolution: {integrity: sha512-+PIZ/6Uk40mphiQJJI1202b+/dYeTVd9ZnMPR80pgiWbjIwvN2zIp4r9et0BgqBuShh48I0gttPlAXA7WVvBxw==}
     engines: {node: '>= 12.0.0'}
     peerDependencies:
@@ -19047,7 +19053,6 @@ packages:
       babel-plugin-macros:
         optional: true
     dependencies:
-      '@babel/core': 7.18.2
       react: 18.0.0
     dev: false
 

test/cache/.gitignore

@@ -0,0 +1 @@
+cache/*

test/cache/package.json

@@ -0,0 +1,11 @@
+{
+  "name": "@vitest/test-cache",
+  "private": true,
+  "scripts": {
+    "test": "vitest",
+    "coverage": "vitest run --coverage"
+  },
+  "devDependencies": {
+    "vitest": "workspace:*"
+  }
+}

test/cache/test/clear-cache.test.ts

@@ -0,0 +1,26 @@
+import fs, { promises as fsp } from 'fs'
+import { resolve } from 'pathe'
+import { describe, expect, test } from 'vitest'
+import { VitestCache } from '../../../packages/vitest/src/node/cache/index'
+
+const root = resolve(__dirname, '..')
+
+const pathBase = resolve(root, 'cache/.vitest-base')
+const pathCustom = resolve(root, 'cache/.vitest-custom')
+
+describe('vitest cache', async () => {
+  await fsp.mkdir(pathBase, { recursive: true })
+  await fsp.mkdir(pathCustom, { recursive: true })
+
+  test('clears cache without specifying config path', async () => {
+    await VitestCache.clearCache({})
+
+    expect(fs.existsSync(pathBase)).toBe(false)
+  })
+
+  test('clears cache with specified config path', async () => {
+    await VitestCache.clearCache({ config: 'vitest-custom.config.ts' })
+
+    expect(fs.existsSync(pathCustom)).toBe(false)
+  })
+})

test/cache/vitest-custom.config.ts

@@ -0,0 +1,9 @@
+import { defineConfig } from 'vite'
+
+export default defineConfig({
+  test: {
+    cache: {
+      dir: 'cache/.vitest-custom',
+    },
+  },
+})

test/cache/vitest.config.ts

@@ -0,0 +1,10 @@
+import { defineConfig } from 'vite'
+
+export default defineConfig({
+  test: {
+    threads: false,
+    cache: {
+      dir: 'cache/.vitest-base',
+    },
+  },
+})

test/core/test/sequelizers.test.ts

@@ -0,0 +1,97 @@
+import type { Vitest } from 'vitest'
+import { describe, expect, test, vi } from 'vitest'
+import { BaseSequelizer } from '../../../packages/vitest/src/node/sequelizers/BaseSequelizer'
+
+const buildCtx = () => {
+  return {
+    state: {
+      getFileTestResults: vi.fn(),
+      getFileStats: vi.fn(),
+    },
+  } as unknown as Vitest
+}
+
+describe('test sequelizers', () => {
+  test('sorting when no info is available', async () => {
+    const sequelizer = new BaseSequelizer(buildCtx())
+    const files = ['a', 'b', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(files)
+  })
+
+  test('prioritaze unknown files', async () => {
+    const ctx = buildCtx()
+    vi.spyOn(ctx.state, 'getFileStats').mockImplementation((file) => {
+      if (file === 'b')
+        return { size: 2 }
+    })
+    const sequelizer = new BaseSequelizer(ctx)
+    const files = ['b', 'a', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(['a', 'c', 'b'])
+  })
+
+  test('sort by size, larger first', async () => {
+    const ctx = buildCtx()
+    vi.spyOn(ctx.state, 'getFileStats').mockImplementation((file) => {
+      if (file === 'a')
+        return { size: 1 }
+      if (file === 'b')
+        return { size: 2 }
+      if (file === 'c')
+        return { size: 3 }
+    })
+    const sequelizer = new BaseSequelizer(ctx)
+    const files = ['b', 'a', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(['c', 'b', 'a'])
+  })
+
+  test('sort by results, failed first', async () => {
+    const ctx = buildCtx()
+    vi.spyOn(ctx.state, 'getFileTestResults').mockImplementation((file) => {
+      if (file === 'a')
+        return { failed: false, duration: 1 }
+      if (file === 'b')
+        return { failed: true, duration: 1 }
+      if (file === 'c')
+        return { failed: true, duration: 1 }
+    })
+    const sequelizer = new BaseSequelizer(ctx)
+    const files = ['b', 'a', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(['b', 'c', 'a'])
+  })
+
+  test('sort by results, long first', async () => {
+    const ctx = buildCtx()
+    vi.spyOn(ctx.state, 'getFileTestResults').mockImplementation((file) => {
+      if (file === 'a')
+        return { failed: true, duration: 1 }
+      if (file === 'b')
+        return { failed: true, duration: 2 }
+      if (file === 'c')
+        return { failed: true, duration: 3 }
+    })
+    const sequelizer = new BaseSequelizer(ctx)
+    const files = ['b', 'a', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(['c', 'b', 'a'])
+  })
+
+  test('sort by results, long and failed first', async () => {
+    const ctx = buildCtx()
+    vi.spyOn(ctx.state, 'getFileTestResults').mockImplementation((file) => {
+      if (file === 'a')
+        return { failed: false, duration: 1 }
+      if (file === 'b')
+        return { failed: false, duration: 6 }
+      if (file === 'c')
+        return { failed: true, duration: 3 }
+    })
+    const sequelizer = new BaseSequelizer(ctx)
+    const files = ['b', 'a', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(['c', 'b', 'a'])
+  })
+})