feat: add environment::listen (#18263)

vitejs · Oct 17, 2024 · 4d5f51d · 4d5f51d
1 parent e59760c
commit 4d5f51d
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 30 deletions.
diff --git a/packages/vite/src/node/server/environment.ts b/packages/vite/src/node/server/environment.ts
@@ -20,6 +20,7 @@ import {
 } from '../optimizer/optimizer'
 import { resolveEnvironmentPlugins } from '../plugin'
 import { ERR_OUTDATED_OPTIMIZED_DEP } from '../constants'
+import type { ViteDevServer } from '../server'
 import { EnvironmentModuleGraph } from './moduleGraph'
 import type { EnvironmentModuleNode } from './moduleGraph'
 import type { HotChannel } from './hmr'
@@ -33,6 +34,7 @@ import {
 } from './pluginContainer'
 import type { RemoteEnvironmentTransport } from './environmentTransport'
 import { isWebSocketServer } from './ws'
+import { warmupFiles } from './warmup'
 
 export interface DevEnvironmentContext {
   hot: false | HotChannel
@@ -149,7 +151,15 @@ export class DevEnvironment extends BaseEnvironment {
     }
   }
 
-  async init(options?: { watcher?: FSWatcher }): Promise<void> {
+  async init(options?: {
+    watcher?: FSWatcher
+    /**
+     * the previous instance used for the environment with the same name
+     *
+     * when using, the consumer should check if it's an instance generated from the same class or factory function
+     */
+    previousInstance?: DevEnvironment
+  }): Promise<void> {
     if (this._initiated) {
       return
     }
@@ -162,6 +172,18 @@ export class DevEnvironment extends BaseEnvironment {
     )
   }
 
+  /**
+   * When the dev server is restarted, the methods are called in the following order:
+   * - new instance `init`
+   * - previous instance `close`
+   * - new instance `listen`
+   */
+  async listen(server: ViteDevServer): Promise<void> {
+    this.hot.listen()
+    await this.depsOptimizer?.init()
+    warmupFiles(server, this)
+  }
+
   fetchModule(
     id: string,
     importer?: string,

diff --git a/packages/vite/src/node/server/index.ts b/packages/vite/src/node/server/index.ts
@@ -93,7 +93,6 @@ import { openBrowser as _openBrowser } from './openBrowser'
 import type { TransformOptions, TransformResult } from './transformRequest'
 import { transformRequest } from './transformRequest'
 import { searchForPackageRoot, searchForWorkspaceRoot } from './searchRoot'
-import { warmupFiles } from './warmup'
 import type { DevEnvironment } from './environment'
 
 export interface ServerOptions extends CommonServerOptions {
@@ -418,12 +417,15 @@ export interface ResolvedServerUrls {
 export function createServer(
   inlineConfig: InlineConfig = {},
 ): Promise<ViteDevServer> {
-  return _createServer(inlineConfig, { hotListen: true })
+  return _createServer(inlineConfig, { listen: true })
 }
 
 export async function _createServer(
   inlineConfig: InlineConfig = {},
-  options: { hotListen: boolean },
+  options: {
+    listen: boolean
+    previousEnvironments?: Record<string, DevEnvironment>
+  },
 ): Promise<ViteDevServer> {
   const config = await resolveConfig(inlineConfig, 'serve')
 
@@ -499,7 +501,8 @@ export async function _createServer(
   }
 
   for (const environment of Object.values(environments)) {
-    await environment.init({ watcher })
+    const previousInstance = options.previousEnvironments?.[environment.name]
+    await environment.init({ watcher, previousInstance })
   }
 
   // Backward compatibility
@@ -902,7 +905,7 @@ export async function _createServer(
   // this code is to avoid calling buildStart multiple times
   let initingServer: Promise<void> | undefined
   let serverInited = false
-  const initServer = async () => {
+  const initServer = async (onListen: boolean) => {
     if (serverInited) return
     if (initingServer) return initingServer
 
@@ -912,14 +915,13 @@ export async function _createServer(
       // buildStart will be called when the first request is transformed
       await environments.client.pluginContainer.buildStart()
 
-      await Promise.all(
-        Object.values(server.environments).map((environment) =>
-          environment.depsOptimizer?.init(),
-        ),
-      )
+      // ensure ws server started
+      if (onListen || options.listen) {
+        await Promise.all(
+          Object.values(environments).map((e) => e.listen(server)),
+        )
+      }
 
-      // TODO: move warmup call inside environment init()
-      warmupFiles(server)
       initingServer = undefined
       serverInited = true
     })()
@@ -931,20 +933,15 @@ export async function _createServer(
     const listen = httpServer.listen.bind(httpServer)
     httpServer.listen = (async (port: number, ...args: any[]) => {
       try {
-        // ensure ws server started
-        Object.values(environments).forEach((e) => e.hot.listen())
-        await initServer()
+        await initServer(true)
       } catch (e) {
         httpServer.emit('error', e)
         return
       }
       return listen(port, ...args)
     }) as any
   } else {
-    if (options.hotListen) {
-      Object.values(environments).forEach((e) => e.hot.listen())
-    }
-    await initServer()
+    await initServer(false)
   }
 
   return server
@@ -1118,7 +1115,10 @@ async function restartServer(server: ViteDevServer) {
     let newServer: ViteDevServer | null = null
     try {
       // delay ws server listen
-      newServer = await _createServer(inlineConfig, { hotListen: false })
+      newServer = await _createServer(inlineConfig, {
+        listen: false,
+        previousEnvironments: server.environments,
+      })
     } catch (err: any) {
       server.config.logger.error(err.message, {
         timestamp: true,
@@ -1151,7 +1151,9 @@ async function restartServer(server: ViteDevServer) {
   if (!middlewareMode) {
     await server.listen(port, true)
   } else {
-    Object.values(server.environments).forEach((e) => e.hot.listen())
+    await Promise.all(
+      Object.values(server.environments).map((e) => e.listen(server)),
+    )
   }
   logger.info('server restarted.', { timestamp: true })
 

diff --git a/packages/vite/src/node/server/warmup.ts b/packages/vite/src/node/server/warmup.ts
@@ -7,15 +7,16 @@ import { normalizePath } from '../utils'
 import type { ViteDevServer } from '../index'
 import type { DevEnvironment } from './environment'
 
-export function warmupFiles(server: ViteDevServer): void {
+export function warmupFiles(
+  server: ViteDevServer,
+  environment: DevEnvironment,
+): void {
   const { root } = server.config
-  for (const environment of Object.values(server.environments)) {
-    mapFiles(environment.config.dev.warmup, root).then((files) => {
-      for (const file of files) {
-        warmupFile(server, environment, file)
-      }
-    })
-  }
+  mapFiles(environment.config.dev.warmup, root).then((files) => {
+    for (const file of files) {
+      warmupFile(server, environment, file)
+    }
+  })
 }
 
 async function warmupFile(