From 2207a68d9362bbe9f9ed5738ddee4e9e9825f0ec Mon Sep 17 00:00:00 2001
From: Elias Papavasileiou <eliaspap@pm.me>
Date: Mon, 10 Jun 2024 18:10:11 +0300
Subject: [PATCH] fix: gracefully shutdown preview server on `SIGTERM` (fix
 #12990) (#17333)

---
 packages/vite/src/node/preview.ts      | 25 ++++++++++++++++++++--
 packages/vite/src/node/server/index.ts | 29 +++++++++++---------------
 packages/vite/src/node/shortcuts.ts    | 10 ++++++---
 packages/vite/src/node/utils.ts        | 16 ++++++++++++++
 4 files changed, 58 insertions(+), 22 deletions(-)

diff --git a/packages/vite/src/node/preview.ts b/packages/vite/src/node/preview.ts
index 4d2e1e645bbcdc..38192abea7626d 100644
--- a/packages/vite/src/node/preview.ts
+++ b/packages/vite/src/node/preview.ts
@@ -25,7 +25,13 @@ import { htmlFallbackMiddleware } from './server/middlewares/htmlFallback'
 import { indexHtmlMiddleware } from './server/middlewares/indexHtml'
 import { notFoundMiddleware } from './server/middlewares/notFound'
 import { proxyMiddleware } from './server/middlewares/proxy'
-import { resolveHostname, resolveServerUrls, shouldServeFile } from './utils'
+import {
+  resolveHostname,
+  resolveServerUrls,
+  setupSIGTERMListener,
+  shouldServeFile,
+  teardownSIGTERMListener,
+} from './utils'
 import { printServerUrls } from './logger'
 import { bindCLIShortcuts } from './shortcuts'
 import type { BindCLIShortcutsOptions } from './shortcuts'
@@ -137,11 +143,16 @@ export async function preview(
   const options = config.preview
   const logger = config.logger
 
+  const closeHttpServer = createServerCloseFn(httpServer)
+
   const server: PreviewServer = {
     config,
     middlewares: app,
     httpServer,
-    close: createServerCloseFn(httpServer),
+    async close() {
+      teardownSIGTERMListener(closeServerAndExit)
+      await closeHttpServer()
+    },
     resolvedUrls: null,
     printUrls() {
       if (server.resolvedUrls) {
@@ -155,6 +166,16 @@ export async function preview(
     },
   }
 
+  const closeServerAndExit = async () => {
+    try {
+      await server.close()
+    } finally {
+      process.exit()
+    }
+  }
+
+  setupSIGTERMListener(closeServerAndExit)
+
   // apply server hooks from plugins
   const postHooks: ((() => void) | void)[] = []
   for (const hook of config.getSortedPluginHooks('configurePreviewServer')) {
diff --git a/packages/vite/src/node/server/index.ts b/packages/vite/src/node/server/index.ts
index 8ba189fd9f728c..9b24de6ede4e2f 100644
--- a/packages/vite/src/node/server/index.ts
+++ b/packages/vite/src/node/server/index.ts
@@ -35,6 +35,8 @@ import {
   promiseWithResolvers,
   resolveHostname,
   resolveServerUrls,
+  setupSIGTERMListener,
+  teardownSIGTERMListener,
 } from '../utils'
 import { getFsUtils } from '../fsUtils'
 import { ssrLoadModule } from '../ssr/ssrModuleLoader'
@@ -502,8 +504,6 @@ export async function _createServer(
   const container = await createPluginContainer(config, moduleGraph, watcher)
   const closeHttpServer = createServerCloseFn(httpServer)
 
-  let exitProcess: () => void
-
   const devHtmlTransformFn = createDevHtmlTransformFn(config)
 
   const onCrawlEndCallbacks: (() => void)[] = []
@@ -638,10 +638,7 @@ export async function _createServer(
     },
     async close() {
       if (!middlewareMode) {
-        process.off('SIGTERM', exitProcess)
-        if (process.env.CI !== 'true') {
-          process.stdin.off('end', exitProcess)
-        }
+        teardownSIGTERMListener(closeServerAndExit)
       }
       await Promise.allSettled([
         watcher.close(),
@@ -736,20 +733,18 @@ export async function _createServer(
     },
   })
 
-  if (!middlewareMode) {
-    exitProcess = async () => {
-      try {
-        await server.close()
-      } finally {
-        process.exit()
-      }
-    }
-    process.once('SIGTERM', exitProcess)
-    if (process.env.CI !== 'true') {
-      process.stdin.on('end', exitProcess)
+  const closeServerAndExit = async () => {
+    try {
+      await server.close()
+    } finally {
+      process.exit()
     }
   }
 
+  if (!middlewareMode) {
+    setupSIGTERMListener(closeServerAndExit)
+  }
+
   const onHMRUpdate = async (
     type: 'create' | 'delete' | 'update',
     file: string,
diff --git a/packages/vite/src/node/shortcuts.ts b/packages/vite/src/node/shortcuts.ts
index 51584c66d48021..1bac571627008b 100644
--- a/packages/vite/src/node/shortcuts.ts
+++ b/packages/vite/src/node/shortcuts.ts
@@ -126,7 +126,11 @@ const BASE_DEV_SHORTCUTS: CLIShortcut<ViteDevServer>[] = [
     key: 'q',
     description: 'quit',
     async action(server) {
-      await server.close().finally(() => process.exit())
+      try {
+        await server.close()
+      } finally {
+        process.exit()
+      }
     },
   },
 ]
@@ -148,9 +152,9 @@ const BASE_PREVIEW_SHORTCUTS: CLIShortcut<PreviewServer>[] = [
   {
     key: 'q',
     description: 'quit',
-    action(server) {
+    async action(server) {
       try {
-        server.httpServer.close()
+        await server.close()
       } finally {
         process.exit()
       }
diff --git a/packages/vite/src/node/utils.ts b/packages/vite/src/node/utils.ts
index 048508a1bad7b5..1064c549bcf069 100644
--- a/packages/vite/src/node/utils.ts
+++ b/packages/vite/src/node/utils.ts
@@ -1438,3 +1438,19 @@ export function partialEncodeURIPath(uri: string): string {
   const postfix = filePath !== uri ? uri.slice(filePath.length) : ''
   return filePath.replaceAll('%', '%25') + postfix
 }
+
+export const setupSIGTERMListener = (callback: () => Promise<void>): void => {
+  process.once('SIGTERM', callback)
+  if (process.env.CI !== 'true') {
+    process.stdin.on('end', callback)
+  }
+}
+
+export const teardownSIGTERMListener = (
+  callback: () => Promise<void>,
+): void => {
+  process.off('SIGTERM', callback)
+  if (process.env.CI !== 'true') {
+    process.stdin.off('end', callback)
+  }
+}