refactor(hmr): keep buffer implementation internal, expose queueUpdate

packages/vite/src/client/client.ts

@@ -30,7 +30,6 @@ const socketHost = `${__HMR_HOSTNAME__ || importMetaUrl.hostname}:${
 }${__HMR_BASE__}`
 const directSocketHost = __HMR_DIRECT_TARGET__
 const base = __BASE__ || '/'
-const messageBuffer: string[] = []
 
 let socket: WebSocket
 try {
@@ -134,38 +133,45 @@ const debounceReload = (time: number) => {
 }
 const pageReload = debounceReload(50)
 
-const hmrClient = new HMRClient(console, async function importUpdatedModule({
-  acceptedPath,
-  timestamp,
-  explicitImportRequired,
-  isWithinCircularImport,
-}) {
-  const [acceptedPathWithoutQuery, query] = acceptedPath.split(`?`)
-  const importPromise = import(
-    /* @vite-ignore */
-    base +
-      acceptedPathWithoutQuery.slice(1) +
-      `?${explicitImportRequired ? 'import&' : ''}t=${timestamp}${
-        query ? `&${query}` : ''
-      }`
-  )
-  if (isWithinCircularImport) {
-    importPromise.catch(() => {
-      console.info(
-        `[hmr] ${acceptedPath} failed to apply HMR as it's within a circular import. Reloading page to reset the execution order. ` +
-          `To debug and break the circular import, you can run \`vite --debug hmr\` to log the circular dependency path if a file change triggered it.`,
-      )
-      pageReload()
-    })
-  }
-  return await importPromise
-})
+const hmrClient = new HMRClient(
+  console,
+  {
+    isReady: () => socket && socket.readyState === 1,
+    send: (message) => socket.send(message),
+  },
+  async function importUpdatedModule({
+    acceptedPath,
+    timestamp,
+    explicitImportRequired,
+    isWithinCircularImport,
+  }) {
+    const [acceptedPathWithoutQuery, query] = acceptedPath.split(`?`)
+    const importPromise = import(
+      /* @vite-ignore */
+      base +
+        acceptedPathWithoutQuery.slice(1) +
+        `?${explicitImportRequired ? 'import&' : ''}t=${timestamp}${
+          query ? `&${query}` : ''
+        }`
+    )
+    if (isWithinCircularImport) {
+      importPromise.catch(() => {
+        console.info(
+          `[hmr] ${acceptedPath} failed to apply HMR as it's within a circular import. Reloading page to reset the execution order. ` +
+            `To debug and break the circular import, you can run \`vite --debug hmr\` to log the circular dependency path if a file change triggered it.`,
+        )
+        pageReload()
+      })
+    }
+    return await importPromise
+  },
+)
 
 async function handleMessage(payload: HMRPayload) {
   switch (payload.type) {
     case 'connected':
       console.debug(`[vite] connected.`)
-      sendMessageBuffer()
+      hmrClient.messenger.flush()
       // proxy(nginx, docker) hmr ws maybe caused timeout,
       // so send ping package let ws keep alive.
       setInterval(() => {
@@ -190,7 +196,7 @@ async function handleMessage(payload: HMRPayload) {
       await Promise.all(
         payload.updates.map(async (update): Promise<void> => {
           if (update.type === 'js-update') {
-            return queueUpdate(hmrClient.fetchUpdate(update))
+            return hmrClient.queueUpdate(update)
           }
 
           // css-update
@@ -306,26 +312,6 @@ function hasErrorOverlay() {
   return document.querySelectorAll(overlayId).length
 }
 
-let pending = false
-let queued: Promise<(() => void) | undefined>[] = []
-
-/**
- * buffer multiple hot updates triggered by the same src change
- * so that they are invoked in the same order they were sent.
- * (otherwise the order may be inconsistent because of the http request round trip)
- */
-async function queueUpdate(p: Promise<(() => void) | undefined>) {
-  queued.push(p)
-  if (!pending) {
-    pending = true
-    await Promise.resolve()
-    pending = false
-    const loading = [...queued]
-    queued = []
-    ;(await Promise.all(loading)).forEach((fn) => fn && fn())
-  }
-}
-
 async function waitForSuccessfulPing(
   socketProtocol: string,
   hostAndPath: string,
@@ -435,22 +421,8 @@ export function removeStyle(id: string): void {
   }
 }
 
-function sendMessageBuffer() {
-  if (socket.readyState === 1) {
-    messageBuffer.forEach((msg) => socket.send(msg))
-    messageBuffer.length = 0
-  }
-}
-
 export function createHotContext(ownerPath: string): ViteHotContext {
-  return new HMRContext(ownerPath, hmrClient, {
-    addBuffer(message) {
-      messageBuffer.push(message)
-    },
-    send() {
-      sendMessageBuffer()
-    },
-  })
+  return new HMRContext(hmrClient, ownerPath)
 }
 
 /**

packages/vite/src/shared/hmr.ts

@@ -15,18 +15,23 @@ interface HotCallback {
   fn: (modules: Array<ModuleNamespace | undefined>) => void
 }
 
-interface Connection {
-  addBuffer(message: string): void
-  send(): unknown
+export interface HMRConnection {
+  /**
+   * Checked before sending messages to the client.
+   */
+  isReady(): boolean
+  /**
+   * Send message to the client.
+   */
+  send(messages: string): void
 }
 
 export class HMRContext implements ViteHotContext {
   private newListeners: CustomListenersMap
 
   constructor(
-    private ownerPath: string,
     private hmrClient: HMRClient,
-    private connection: Connection,
+    private ownerPath: string,
   ) {
     if (!hmrClient.dataMap.has(ownerPath)) {
       hmrClient.dataMap.set(ownerPath, {})
@@ -141,8 +146,9 @@ export class HMRContext implements ViteHotContext {
   }
 
   send<T extends string>(event: T, data?: InferCustomEventPayload<T>): void {
-    this.connection.addBuffer(JSON.stringify({ type: 'custom', event, data }))
-    this.connection.send()
+    this.hmrClient.messenger.send(
+      JSON.stringify({ type: 'custom', event, data }),
+    )
   }
 
   private acceptDeps(
@@ -161,6 +167,24 @@ export class HMRContext implements ViteHotContext {
   }
 }
 
+class HMRMessenger {
+  constructor(private connection: HMRConnection) {}
+
+  private queue: string[] = []
+
+  public send(message: string): void {
+    this.queue.push(message)
+    this.flush()
+  }
+
+  public flush(): void {
+    if (this.connection.isReady()) {
+      this.queue.forEach((msg) => this.connection.send(msg))
+      this.queue = []
+    }
+  }
+}
+
 export class HMRClient {
   public hotModulesMap = new Map<string, HotModule>()
   public disposeMap = new Map<string, (data: any) => void | Promise<void>>()
@@ -169,11 +193,16 @@ export class HMRClient {
   public customListenersMap: CustomListenersMap = new Map()
   public ctxToListenersMap = new Map<string, CustomListenersMap>()
 
+  public messenger: HMRMessenger
+
   constructor(
     public logger: Console,
-    // this allows up to implement reloading via different methods depending on the environment
+    connection: HMRConnection,
+    // This allows implementing reloading via different methods depending on the environment
     private importUpdatedModule: (update: Update) => Promise<ModuleNamespace>,
-  ) {}
+  ) {
+    this.messenger = new HMRMessenger(connection)
+  }
 
   public async notifyListeners<T extends string>(
     event: T,
@@ -210,6 +239,26 @@ export class HMRClient {
     )
   }
 
+  private updateQueue: Promise<(() => void) | undefined>[] = []
+  private pendingUpdateQueue = false
+
+  /**
+   * buffer multiple hot updates triggered by the same src change
+   * so that they are invoked in the same order they were sent.
+   * (otherwise the order may be inconsistent because of the http request round trip)
+   */
+  public async queueUpdate(payload: Update): Promise<void> {
+    this.updateQueue.push(this.fetchUpdate(payload))
+    if (!this.pendingUpdateQueue) {
+      this.pendingUpdateQueue = true
+      await Promise.resolve()
+      this.pendingUpdateQueue = false
+      const loading = [...this.updateQueue]
+      this.updateQueue = []
+      ;(await Promise.all(loading)).forEach((fn) => fn && fn())
+    }
+  }
+
   public async fetchUpdate(update: Update): Promise<(() => void) | undefined> {
     const { path, acceptedPath } = update
     const mod = this.hotModulesMap.get(path)