From ff6efb1c4a919f3cace4c0ab2ac69a4f181d2bd8 Mon Sep 17 00:00:00 2001
From: Dominik G <dominik.goepel@gmx.de>
Date: Wed, 29 Dec 2021 05:23:00 +0100
Subject: [PATCH] feat(cli): graceful close of vitest in watch mode (#349)

---
 packages/vitest/src/node/cli.ts   | 31 +++++++++++++++++++++++++++----
 packages/vitest/src/node/index.ts | 15 +++++++++++++--
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/packages/vitest/src/node/cli.ts b/packages/vitest/src/node/cli.ts
index 42494f44efab..f350bbae08ae 100644
--- 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,16 +107,37 @@ 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)
@@ -122,7 +145,7 @@ function registerConsoleShortcuts(ctx: Vitest) {
 
     // 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
index b27c732d9b7a..56621b366d3a 100644
--- 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]>) {