feat(cli): graceful close of vitest in watch mode (#349)

vitest-dev · Dec 29, 2021 · ff6efb1 · ff6efb1
1 parent e9dc52f
commit ff6efb1
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 6 deletions.
diff --git a/packages/vitest/src/node/cli.ts b/packages/vitest/src/node/cli.ts
@@ -7,6 +7,8 @@ import { ensurePackageInstalled } from '../utils'
 import type { Vitest } from './index'
 import { createVitest } from './index'
 
+const CLOSE_TIMEOUT = 1_000
+
 const cli = cac('vitest')
 
 cli
@@ -105,24 +107,45 @@ async function run(cliFilters: string[], options: UserConfig) {
 
   if (!ctx.config.watch) {
     // force process exit if it hangs
-    setTimeout(() => process.exit(), 1000).unref()
+    setTimeout(() => process.exit(), CLOSE_TIMEOUT).unref()
   }
 }
 
+function closeServerAndExitProcess(ctx: Vitest) {
+  const closePromise = ctx.close()
+  let timeout: NodeJS.Timeout
+  const timeoutPromise = new Promise((resolve, reject) => {
+    timeout = setTimeout(() => reject(new Error(`close timed out after ${CLOSE_TIMEOUT}ms`)), CLOSE_TIMEOUT)
+  })
+  Promise.race([closePromise, timeoutPromise]).then(
+    () => {
+      clearTimeout(timeout)
+      process.exit(0)
+    },
+    (err) => {
+      clearTimeout(timeout)
+      console.error('error during close', err)
+      process.exit(1)
+    },
+  )
+}
+
 function registerConsoleShortcuts(ctx: Vitest) {
   readline.emitKeypressEvents(process.stdin)
   process.stdin.setRawMode(true)
   process.stdin.on('keypress', (str: string, key: any) => {
-    if (str === '\x03' || str === '\x1B' || (key && key.ctrl && key.name === 'c')) // ctrl-c or esc
-      process.exit()
+    if (str === '\x03' || str === '\x1B' || (key && key.ctrl && key.name === 'c')) { // ctrl-c or esc
+      closeServerAndExitProcess(ctx)
+      return
+    }
 
     // is running, ignore keypress
     if (ctx.runningPromise)
       return
 
     // press any key to exit on first run
     if (ctx.isFirstRun)
-      process.exit()
+      closeServerAndExitProcess(ctx)
 
     // TODO: add more commands
   })

diff --git a/packages/vitest/src/node/index.ts b/packages/vitest/src/node/index.ts
@@ -37,6 +37,8 @@ class Vitest {
   changedTests: Set<string> = new Set()
   visitedFilesMap: Map<string, RawSourceMap> = new Map()
   runningPromise?: Promise<void>
+  closingPromise?: Promise<void>
+
   isFirstRun = true
 
   restartsCount = 0
@@ -248,8 +250,17 @@ class Vitest {
   }
 
   async close() {
-    await this.pool?.close()
-    await this.server.close()
+    if (!this.closingPromise) {
+      this.closingPromise = Promise.allSettled([
+        this.pool?.close(),
+        this.server.close(),
+      ].filter(Boolean)).then((results) => {
+        results.filter(r => r.status === 'rejected').forEach((err) => {
+          this.error('error during close', (err as PromiseRejectedResult).reason)
+        })
+      })
+    }
+    return this.closingPromise
   }
 
   async report<T extends keyof Reporter>(name: T, ...args: ArgumentsType<Reporter[T]>) {