feat: support multiple HMR clients on the server (#15340)

Author: sheremet-va

docs/guide/api-plugin.md

@@ -423,11 +423,11 @@ Vite plugins can also provide hooks that serve Vite-specific purposes. These hoo
 
   - Filter and narrow down the affected module list so that the HMR is more accurate.
 
-  - Return an empty array and perform complete custom HMR handling by sending custom events to the client:
+  - Return an empty array and perform complete custom HMR handling by sending custom events to the client (example uses `server.hot` which was introduced in Vite 5.1, it is recommended to also use `server.ws` if you support lower versions):
 
     ```js
     handleHotUpdate({ server }) {
-      server.ws.send({
+      server.hot.send({
         type: 'custom',
         event: 'special-update',
         data: {}
@@ -534,7 +534,7 @@ Since Vite 2.9, we provide some utilities for plugins to help handle the communi
 
 ### Server to Client
 
-On the plugin side, we could use `server.ws.send` to broadcast events to all the clients:
+On the plugin side, we could use `server.hot.send` (since Vite 5.1) or `server.ws.send` to broadcast events to all the clients:
 
 ```js
 // vite.config.js
@@ -544,8 +544,8 @@ export default defineConfig({
       // ...
       configureServer(server) {
         // Example: wait for a client to connect before sending a message
-        server.ws.on('connection', () => {
-          server.ws.send('my:greetings', { msg: 'hello' })
+        server.hot.on('connection', () => {
+          server.hot.send('my:greetings', { msg: 'hello' })
         })
       },
     },
@@ -579,7 +579,7 @@ if (import.meta.hot) {
 }
 ```
 
-Then use `server.ws.on` and listen to the events on the server side:
+Then use `server.hot.on` (since Vite 5.1) or `server.ws.on` and listen to the events on the server side:
 
 ```js
 // vite.config.js
@@ -588,7 +588,7 @@ export default defineConfig({
     {
       // ...
       configureServer(server) {
-        server.ws.on('my:from-client', (data, client) => {
+        server.hot.on('my:from-client', (data, client) => {
           console.log('Message from client:', data.msg) // Hey!
           // reply only to the client (if needed)
           client.send('my:ack', { msg: 'Hi! I got your message!' })

packages/vite/src/node/optimizer/optimizer.ts

@@ -473,7 +473,7 @@ async function createDepsOptimizer(
       // reloaded.
       server.moduleGraph.invalidateAll()
 
-      server.ws.send({
+      server.hot.send({
         type: 'full-reload',
         path: '*',
       })

packages/vite/src/node/plugin.ts

@@ -129,7 +129,7 @@ export interface Plugin<A = any> extends RollupPlugin<A> {
    *   the descriptors.
    *
    * - The hook can also return an empty array and then perform custom updates
-   *   by sending a custom hmr payload via server.ws.send().
+   *   by sending a custom hmr payload via server.hot.send().
    *
    * - If the hook doesn't return a value, the hmr update will be performed as
    *   normal.

packages/vite/src/node/plugins/esbuild.ts

@@ -491,7 +491,7 @@ async function reloadOnTsconfigChange(changedFile: string) {
     // server may not be available if vite config is updated at the same time
     if (server) {
       // force full reload
-      server.ws.send({
+      server.hot.send({
         type: 'full-reload',
         path: '*',
       })

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

@@ -2,7 +2,7 @@ import fsp from 'node:fs/promises'
 import path from 'node:path'
 import type { Server } from 'node:http'
 import colors from 'picocolors'
-import type { Update } from 'types/hmrPayload'
+import type { CustomPayload, HMRPayload, Update } from 'types/hmrPayload'
 import type { RollupError } from 'rollup'
 import { CLIENT_DIR } from '../constants'
 import {
@@ -12,7 +12,7 @@ import {
   withTrailingSlash,
   wrapId,
 } from '../utils'
-import type { ViteDevServer } from '..'
+import type { InferCustomEventPayload, ViteDevServer } from '..'
 import { isCSSRequest } from '../plugins/css'
 import { getAffectedGlobModules } from '../plugins/importMetaGlob'
 import { isExplicitImportRequired } from '../plugins/importAnalysis'
@@ -35,6 +35,8 @@ export interface HmrOptions {
   timeout?: number
   overlay?: boolean
   server?: Server
+  /** @internal */
+  channels?: HMRChannel[]
 }
 
 export interface HmrContext {
@@ -51,6 +53,68 @@ interface PropagationBoundary {
   isWithinCircularImport: boolean
 }
 
+export interface HMRBroadcasterClient {
+  /**
+   * Send event to the client
+   */
+  send(payload: HMRPayload): void
+  /**
+   * Send custom event
+   */
+  send(event: string, payload?: CustomPayload['data']): void
+}
+
+export interface HMRChannel {
+  /**
+   * Unique channel name
+   */
+  name: string
+  /**
+   * Broadcast events to all clients
+   */
+  send(payload: HMRPayload): void
+  /**
+   * Send custom event
+   */
+  send<T extends string>(event: T, payload?: InferCustomEventPayload<T>): void
+  /**
+   * Handle custom event emitted by `import.meta.hot.send`
+   */
+  on<T extends string>(
+    event: T,
+    listener: (
+      data: InferCustomEventPayload<T>,
+      client: HMRBroadcasterClient,
+      ...args: any[]
+    ) => void,
+  ): void
+  on(event: 'connection', listener: () => void): void
+  /**
+   * 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
+}
+
+export interface HMRBroadcaster extends Omit<HMRChannel, 'close' | 'name'> {
+  /**
+   * All registered channels. Always has websocket channel.
+   */
+  readonly channels: HMRChannel[]
+  /**
+   * Add a new third-party channel.
+   */
+  addChannel(connection: HMRChannel): HMRBroadcaster
+  close(): Promise<unknown[]>
+}
+
 export function getShortName(file: string, root: string): string {
   return file.startsWith(withTrailingSlash(root))
     ? path.posix.relative(root, file)
@@ -62,7 +126,7 @@ export async function handleHMRUpdate(
   server: ViteDevServer,
   configOnly: boolean,
 ): Promise<void> {
-  const { ws, config, moduleGraph } = server
+  const { hot, config, moduleGraph } = server
   const shortFile = getShortName(file, config.root)
 
   const isConfig = file === config.configFile
@@ -98,7 +162,7 @@ export async function handleHMRUpdate(
 
   // (dev only) the client itself cannot be hot updated.
   if (file.startsWith(withTrailingSlash(normalizedClientDir))) {
-    ws.send({
+    hot.send({
       type: 'full-reload',
       path: '*',
     })
@@ -131,7 +195,7 @@ export async function handleHMRUpdate(
         clear: true,
         timestamp: true,
       })
-      ws.send({
+      hot.send({
         type: 'full-reload',
         path: config.server.middlewareMode
           ? '*'
@@ -153,7 +217,7 @@ export function updateModules(
   file: string,
   modules: ModuleNode[],
   timestamp: number,
-  { config, ws, moduleGraph }: ViteDevServer,
+  { config, hot, moduleGraph }: ViteDevServer,
   afterInvalidation?: boolean,
 ): void {
   const updates: Update[] = []
@@ -202,7 +266,7 @@ export function updateModules(
       colors.green(`page reload `) + colors.dim(file) + reason,
       { clear: !afterInvalidation, timestamp: true },
     )
-    ws.send({
+    hot.send({
       type: 'full-reload',
     })
     return
@@ -218,7 +282,7 @@ export function updateModules(
       colors.dim([...new Set(updates.map((u) => u.path))].join(', ')),
     { clear: !afterInvalidation, timestamp: true },
   )
-  ws.send({
+  hot.send({
     type: 'update',
     updates,
   })
@@ -453,7 +517,7 @@ function isNodeWithinCircularImports(
 
 export function handlePrunedModules(
   mods: Set<ModuleNode>,
-  { ws }: ViteDevServer,
+  { hot }: ViteDevServer,
 ): void {
   // update the disposed modules' hmr timestamp
   // since if it's re-imported, it should re-apply side effects
@@ -463,7 +527,7 @@ export function handlePrunedModules(
     mod.lastHMRTimestamp = t
     debugHmr?.(`[dispose] ${colors.dim(mod.file)}`)
   })
-  ws.send({
+  hot.send({
     type: 'prune',
     paths: [...mods].map((m) => m.url),
   })
@@ -638,3 +702,52 @@ async function readModifiedFile(file: string): Promise<string> {
     return content
   }
 }
+
+export function createHMRBroadcaster(): HMRBroadcaster {
+  const channels: HMRChannel[] = []
+  const readyChannels = new WeakSet<HMRChannel>()
+  const broadcaster: HMRBroadcaster = {
+    get channels() {
+      return [...channels]
+    },
+    addChannel(channel) {
+      if (channels.some((c) => c.name === channel.name)) {
+        throw new Error(`HMR channel "${channel.name}" is already defined.`)
+      }
+      channels.push(channel)
+      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
+    },
+    off(event, listener) {
+      channels.forEach((channel) => channel.off(event, listener))
+      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()))
+    },
+  }
+  return broadcaster
+}