feat: add channels option

packages/vite/src/node/server/hmr.ts

@@ -35,6 +35,8 @@ export interface HmrOptions {
   timeout?: number
   overlay?: boolean
   server?: Server
+  /** @internal */
+  channels?: HMRChannel[]
 }
 
 export interface HmrContext {
@@ -91,10 +93,14 @@ export interface HMRChannel {
    * Unregister event listener
    */
   off(event: string, listener: Function): void
+  /**
+   * Start listening for messages
+   */
+  listen(): void
   /**
    * Disconnect all clients, called when server is closed or restarted.
    */
-  close(): void | Promise<void>
+  close(): void
 }
 
 export interface HMRBroadcaster extends Omit<HMRChannel, 'close' | 'name'> {
@@ -705,6 +711,7 @@ async function readModifiedFile(file: string): Promise<string> {
 
 export function createHMRBroadcaster(): HMRBroadcaster {
   const channels: HMRChannel[] = []
+  const readyChannels = new WeakSet<HMRChannel>()
   const broadcaster: HMRBroadcaster = {
     get channels() {
       return [...channels]
@@ -717,16 +724,33 @@ export function createHMRBroadcaster(): HMRBroadcaster {
       return broadcaster
     },
     on(event: string, listener: (...args: any[]) => any) {
+      // emit connection event only when all channels are ready
+      if (event === 'connection') {
+        // make a copy so we don't wait for channels that might be added after this is triggered
+        const channels = this.channels
+        channels.forEach((channel) =>
+          channel.on('connection', () => {
+            readyChannels.add(channel)
+            if (channels.every((c) => readyChannels.has(c))) {
+              listener()
+            }
+          }),
+        )
+        return
+      }
       channels.forEach((channel) => channel.on(event, listener))
-      return broadcaster
+      return
     },
     off(event, listener) {
       channels.forEach((channel) => channel.off(event, listener))
-      return broadcaster
+      return
     },
     send(...args: any[]) {
       channels.forEach((channel) => channel.send(...(args as [any])))
     },
+    listen() {
+      channels.forEach((channel) => channel.listen())
+    },
     close() {
       return Promise.all(channels.map((channel) => channel.close()))
     },

packages/vite/src/node/server/index.ts

@@ -239,7 +239,7 @@ export interface ViteDevServer {
    * HMR broadcaster that can be used to send custom HMR messages to the client
    *
    * Always sends a message to at least a WebSocket client. Any third party can
-   * add a channel to the broadcaster to process messages so be aware of that
+   * add a channel to the broadcaster to process messages
    */
   hot: HMRBroadcaster
   /**
@@ -387,12 +387,12 @@ export interface ResolvedServerUrls {
 export function createServer(
   inlineConfig: InlineConfig = {},
 ): Promise<ViteDevServer> {
-  return _createServer(inlineConfig, { ws: true })
+  return _createServer(inlineConfig, { hotListen: true })
 }
 
 export async function _createServer(
   inlineConfig: InlineConfig = {},
-  options: { ws: boolean },
+  options: { hotListen: boolean },
 ): Promise<ViteDevServer> {
   const config = await resolveConfig(inlineConfig, 'serve')
 
@@ -414,6 +414,9 @@ export async function _createServer(
 
   const ws = createWebSocketServer(httpServer, config, httpsOptions)
   const hot = createHMRBroadcaster().addChannel(ws)
+  if (typeof config.server.hmr === 'object' && config.server.hmr.channels) {
+    config.server.hmr.channels.forEach((channel) => hot.addChannel(channel))
+  }
 
   if (httpServer) {
     setClientErrorHandler(httpServer, config.logger)
@@ -836,7 +839,7 @@ export async function _createServer(
     httpServer.listen = (async (port: number, ...args: any[]) => {
       try {
         // ensure ws server started
-        ws.listen()
+        hot.listen()
         await initServer()
       } catch (e) {
         httpServer.emit('error', e)
@@ -845,8 +848,8 @@ export async function _createServer(
       return listen(port, ...args)
     }) as any
   } else {
-    if (options.ws) {
-      ws.listen()
+    if (options.hotListen) {
+      hot.listen()
     }
     await initServer()
   }
@@ -997,7 +1000,7 @@ async function restartServer(server: ViteDevServer) {
     let newServer = null
     try {
       // delay ws server listen
-      newServer = await _createServer(inlineConfig, { ws: false })
+      newServer = await _createServer(inlineConfig, { hotListen: false })
     } catch (err: any) {
       server.config.logger.error(err.message, {
         timestamp: true,
@@ -1030,7 +1033,7 @@ async function restartServer(server: ViteDevServer) {
   if (!middlewareMode) {
     await server.listen(port, true)
   } else {
-    server.ws.listen()
+    server.hot.listen()
   }
   logger.info('server restarted.', { timestamp: true })